diff --git a/.github/workflows/docs.yaml b/.github/workflows/docs.yaml index 85c49fadbbf4..1be69475745c 100644 --- a/.github/workflows/docs.yaml +++ b/.github/workflows/docs.yaml @@ -94,10 +94,11 @@ jobs: working-directory: ./docs env: # "deploy" branches build the full set of versions so every page - # has a complete version dropdown: main, develop, tags >= v2.0.0. - # v1.x tags and release/ branches are excluded to keep it lean and mean. + # has a complete version dropdown: main, develop, tags >= v2.0.0 + # (including pre-release suffixes like -beta or -rc1). v1.x tags and + # release/ branches are excluded. SMV_BRANCH_WHITELIST: '^(main|develop)$' - SMV_TAG_WHITELIST: '^v[2-9]\d*\.\d+\.\d+$' + SMV_TAG_WHITELIST: '^v[2-9]\d*\.\d+\.\d+(-[A-Za-z0-9.]+)?$' run: | git fetch --prune --unshallow --tags git checkout --detach HEAD diff --git a/.gitignore b/.gitignore index 5d4c8f954fc0..60989ad5dd1b 100644 --- a/.gitignore +++ b/.gitignore @@ -79,3 +79,12 @@ _build # Ruff cache **/.ruff_cache/ + +# Dev-time files, generated stuff +**/__* + +# Isaac Lab CI environments in native mode +**/_isaaclab_install_ci_* + +# Superpowers (Claude Code plugin artifacts) +docs/superpowers/ diff --git a/AGENTS.md b/AGENTS.md index 70c1d15158ab..91f858751f50 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -64,7 +64,7 @@ We use a wrapped python call within `./isaaclab.sh`. ### Pre-commit (lint/format hooks) -**CRITICAL: Always run pre-commit hooks BEFORE committing, not after.** +**CRITICAL: Always run pre-commit hooks BEFORE committing and BEFORE pushing.** Proper workflow: 1. Make your code changes @@ -73,15 +73,17 @@ Proper workflow: 4. Stage the modified files with `git add` 5. Run `./isaaclab.sh -f` again to ensure all checks pass 6. Only then create your commit with `git commit` +7. Verify pre-commit still passes before pushing — never push commits that haven't been checked ```bash # Run pre-commit checks on all files ./isaaclab.sh -f ``` -**Common mistake to avoid:** +**Common mistakes to avoid:** - Don't commit first and then run pre-commit (requires amending commits) -- Do run pre-commit before committing (clean workflow) +- Don't push before running pre-commit (pushes broken code to the remote) +- Do run pre-commit before committing and before pushing (clean workflow) **When reviewing code** (e.g. via a code-reviewer agent), always run `./isaaclab.sh -f` as part of the review to catch formatting or lint issues early. @@ -152,6 +154,7 @@ Follow conventional commit message practices. ## Sandbox & Networking - Network access (e.g., `git push`) is blocked by the sandbox. Use `dangerouslyDisableSandbox: true` so the user gets an approval prompt — don't ask them to run it manually. +- **Never push to `origin` (`isaac-sim/IsaacLab`).** The `origin` remote is the public upstream repository. Push to your own fork remote (e.g., `antoine`, `alex`) or to the remote of the PR you are working on. If the correct remote is unclear, ask the user before pushing. ## GitHub Actions and CI/CD diff --git a/docs/conf.py b/docs/conf.py index fcd2bcb9eca0..510fa3bc9280 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -186,6 +186,7 @@ "omni.timeline", "omni.ui", "gym", + "gymnasium", "skrl", "stable_baselines3", "rsl_rl", @@ -304,8 +305,10 @@ smv_remote_whitelist = r"^.*$" # Whitelist pattern for branches (set to None to ignore all branches) smv_branch_whitelist = os.getenv("SMV_BRANCH_WHITELIST", r"^(main|develop|release/.*)$") -# Whitelist pattern for tags (set to None to ignore all tags) -smv_tag_whitelist = os.getenv("SMV_TAG_WHITELIST", r"^v[1-9]\d*\.\d+\.\d+$") +# Whitelist pattern for tags (set to None to ignore all tags). +# Matches vMAJOR.MINOR.PATCH with an optional pre-release suffix like -beta or -rc1, +# so tags like v3.0.0-beta show up in the version selector. +smv_tag_whitelist = os.getenv("SMV_TAG_WHITELIST", r"^v[1-9]\d*\.\d+\.\d+(-[A-Za-z0-9.]+)?$") html_sidebars = { "**": ["navbar-logo.html", "versioning.html", "icon-links.html", "search-field.html", "sbt-sidebar-nav.html"] } diff --git a/docs/index.rst b/docs/index.rst index 9a8d91664f7b..6c8e6ba9ac41 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -52,7 +52,7 @@ For more information about the framework, please refer to the `technical report License -======= +======== The Isaac Lab framework is open-sourced under the BSD-3-Clause license, with certain parts under Apache-2.0 license. Please refer to :ref:`license` for more details. 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/features/isaac_teleop.rst b/docs/source/features/isaac_teleop.rst index ac910d6ed4ed..733150259586 100644 --- a/docs/source/features/isaac_teleop.rst +++ b/docs/source/features/isaac_teleop.rst @@ -115,8 +115,10 @@ and Isaac Lab. It composes three collaborators: Isaac Sim's XR bridge, creates the ``TeleopSession``, and steps it each frame to produce an action tensor. -* **CommandHandler** -- registers and dispatches START / STOP / RESET callbacks triggered by XR UI - buttons or the message bus. +* **CommandHandler** -- lightweight callback registry for START / STOP / RESET commands. Scripts + can register callbacks via :meth:`~isaaclab_teleop.IsaacTeleopDevice.add_callback`, but the + primary control path uses :func:`~isaaclab_teleop.poll_control_events` (see + :ref:`isaac-teleop-control-states`). .. dropdown:: Session lifecycle details @@ -127,6 +129,104 @@ and Isaac Lab. It composes three collaborators: the session is not yet ready or has been torn down. +.. _isaac-teleop-control-states: + +Teleop Control States (Start / Stop / Reset) +--------------------------------------------- + +Isaac Lab supports remote teleop control commands -- **start**, **stop**, and **reset** -- sent +from the XR headset to the simulation. These are used to begin and end demonstration recording, +pause the robot, or reset the environment without touching the simulation host. + +How it works +~~~~~~~~~~~~ + +By default, every :class:`~isaaclab_teleop.IsaacTeleopCfg` enables a control message channel +using the well-known UUID ``uuid5(NAMESPACE_DNS, "teleop_command")``. The channel is created as +a ``teleop_control_pipeline`` inside TeleopCore's :class:`TeleopSession`, which means: + +1. A :class:`~isaacteleop.retargeting_engine.deviceio_source_nodes.MessageChannelSource` opens an + OpenXR opaque data channel (``XR_NV_opaque_data_channel``) with the agreed-upon UUID. +2. The CloudXR JS client (or any other client) discovers the channel by UUID and sends UTF-8 + JSON commands:: + + {"type": "teleop_command", "message": {"command": "start teleop"}} + {"type": "teleop_command", "message": {"command": "stop teleop"}} + {"type": "teleop_command", "message": {"command": "reset teleop"}} + +3. A :class:`~isaaclab_teleop.teleop_message_processor.TeleopMessageProcessor` parses these + payloads and produces boolean pulse signals (``run_toggle``, ``kill``, ``reset``). +4. :class:`~isaacteleop.teleop_session_manager.DefaultTeleopStateManager` consumes the + boolean signals, runs its state machine (edge detection, fail-safe), and produces + ``teleop_state`` (one-hot) and ``reset_event`` (bool pulse) outputs. +5. TeleopCore decodes these outputs into ``ExecutionEvents`` and injects them into every + retargeter's ``ComputeContext``, so stateful retargeters can react to state changes + (e.g. reinitializing cross-step state on reset). + +Polling control events in your script +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +Use :func:`~isaaclab_teleop.poll_control_events` to read the latest control state each frame: + +.. code-block:: python + + from isaaclab_teleop import poll_control_events + + with IsaacTeleopDevice(cfg) as device: + running = False + while sim_app.is_running(): + action = device.advance() + + ctrl = poll_control_events(device) + if ctrl.is_active is not None: + running = ctrl.is_active # True after "start", False after "stop" + if ctrl.should_reset: + env.reset() # "reset" command received this frame + + if action is not None and running: + env.step(action.repeat(num_envs, 1)) + else: + env.sim.render() + +:class:`~isaaclab_teleop.ControlEvents` has two fields: + +* ``is_active`` -- ``True`` after a "start" command, ``False`` after "stop", ``None`` when no + command has been received yet (callers should leave their own flag unchanged). +* ``should_reset`` -- ``True`` for exactly one frame after a "reset" command. + +Disabling the control channel +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +If you do not need headset-driven start/stop/reset (e.g. keyboard-only workflows), set +``control_channel_uuid=None`` in your config: + +.. code-block:: python + + IsaacTeleopCfg( + pipeline_builder=_build_my_pipeline, + control_channel_uuid=None, # no opaque data channel created + ) + +Using a custom channel UUID +~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +To use a different channel UUID (e.g. for a separate control protocol), pass any 16-byte +``bytes`` value: + +.. code-block:: python + + import uuid + + MY_UUID = uuid.uuid5(uuid.NAMESPACE_DNS, "my_custom_control").bytes + + IsaacTeleopCfg( + pipeline_builder=_build_my_pipeline, + control_channel_uuid=MY_UUID, + ) + +The CloudXR JS client must be updated to discover this UUID when sending commands. + + .. _isaac-teleop-retargeting: Retargeting Framework @@ -908,6 +1008,10 @@ See the :ref:`isaaclab_teleop-api` for full class and function documentation: * :class:`~isaaclab_teleop.IsaacTeleopCfg` * :class:`~isaaclab_teleop.IsaacTeleopDevice` * :func:`~isaaclab_teleop.create_isaac_teleop_device` +* :class:`~isaaclab_teleop.ControlEvents` +* :class:`~isaaclab_teleop.SupportsControlEvents` +* :func:`~isaaclab_teleop.poll_control_events` +* :data:`~isaaclab_teleop.TELEOP_CONTROL_CHANNEL_UUID` * :class:`~isaaclab_teleop.XrCfg` * :class:`~isaaclab_teleop.XrAnchorRotationMode` diff --git a/docs/source/features/visualization.rst b/docs/source/features/visualization.rst index 13d573308efb..b9d23a45cf91 100644 --- a/docs/source/features/visualization.rst +++ b/docs/source/features/visualization.rst @@ -104,17 +104,14 @@ You can also configure custom visualizers in the code by defining ``VisualizerCf sim_cfg = SimulationCfg( visualizer_cfgs=[ KitVisualizerCfg( - viewport_name="Visualizer Viewport", - create_viewport=True, - dock_position="SAME", - window_width=1280, - window_height=720, - camera_position=(0.0, 0.0, 20.0), # high top down view - camera_target=(0.0, 0.0, 0.0), + # Omit create_viewport (default False) to use the active viewport; set + # create_viewport=True and optionally viewport_name to add a dedicated window. + eye=(0.0, 0.0, 20.0), # high top down view + lookat=(0.0, 0.0, 0.0), ), NewtonVisualizerCfg( - camera_position=(5.0, 5.0, 5.0), # closer quarter view - camera_target=(0.0, 0.0, 0.0), + eye=(5.0, 5.0, 5.0), # closer quarter view + lookat=(0.0, 0.0, 0.0), show_joints=True, ), RerunVisualizerCfg( @@ -142,6 +139,21 @@ The effective visualizer mode is resolved from both CLI and ``SimulationCfg.visu For the migration-focused summary and deprecation context, see :doc:`/source/migration/migrating_to_isaaclab_3-0`. +Partial Visualization +~~~~~~~~~~~~~~~~~~~~~ + +Visualizers can be configured to visualize just a subset of environments. +This is called partial visualization. + +There are 3 fields exposed in the ``VisualizerCfg`` for selecting environments for partial visualization: + +- ``max_visible_envs`` caps how many envs are shown. +- ``visible_env_indices`` explicitly selects the envs to visualize. +- ``randomly_sample_visible_envs`` (default ``True``): when ``visible_env_indices`` is unset and ``max_visible_envs`` is set, + enables randomly sampling the selected envs. If disabled, the first ``max_visible_envs`` envs are selected. + +Also, there is a CLI arg ``--max_visible_envs`` that overrides ``VisualizerCfg.max_visible_envs`` for the run. + .. _visualization-common-modes: .. list-table:: Common modes @@ -193,20 +205,18 @@ Omniverse Visualizer from isaaclab_visualizers.kit import KitVisualizerCfg visualizer_cfg = KitVisualizerCfg( - # Viewport settings - viewport_name="Visualizer Viewport", # Viewport window name - create_viewport=True, # Create new viewport vs. use existing - dock_position="SAME", # Docking: 'LEFT', 'RIGHT', 'BOTTOM', 'SAME' - window_width=1280, # Viewport width in pixels - window_height=720, # Viewport height in pixels - - # Camera settings - camera_position=(8.0, 8.0, 3.0), # Initial camera position (x, y, z) - camera_target=(0.0, 0.0, 0.0), # Camera look-at target - - # Feature toggles - enable_markers=True, # Enable visualization markers - enable_live_plots=True, # Enable live plots (auto-expands frames) + # Viewport: default is create_viewport=False (use active viewport). + # Set create_viewport=True to create a docked window; viewport_name=None uses the default name. + create_viewport=False, + dock_position="SAME", + window_width=1280, + window_height=720, + + eye=(8.0, 8.0, 3.0), + lookat=(0.0, 0.0, 0.0), + + enable_markers=True, + enable_live_plots=True, ) @@ -217,7 +227,7 @@ Newton Visualizer - Lightweight OpenGL rendering with low overhead - Visualization markers (joints, contacts, springs, COM) -- Training and rendering pause controls +- Simulation and rendering pause controls - Adjustable update frequency for performance tuning - Some customizable rendering options (shadows, sky, wireframe) @@ -255,8 +265,8 @@ Newton Visualizer window_height=1080, # Window height in pixels # Camera settings - camera_position=(8.0, 8.0, 3.0), # Initial camera position (x, y, z) - camera_target=(0.0, 0.0, 0.0), # Camera look-at target + eye=(8.0, 8.0, 3.0), # Initial camera position (x, y, z) + lookat=(0.0, 0.0, 0.0), # Camera look-at target # Performance tuning update_frequency=1, # Update every N frames (1=every frame) @@ -303,8 +313,8 @@ Rerun Visualizer bind_address="0.0.0.0", # Endpoint host formatting/reuse checks # Camera settings - camera_position=(8.0, 8.0, 3.0), # Initial camera position (x, y, z) - camera_target=(0.0, 0.0, 0.0), # Camera look-at target + eye=(8.0, 8.0, 3.0), # Initial camera position (x, y, z) + lookat=(0.0, 0.0, 0.0), # Camera look-at target # History settings keep_historical_data=False, # Keep transforms for time scrubbing @@ -350,7 +360,7 @@ server, allowing you to view and interact with the scene from any browser. open_browser=True, label="Isaac Lab Simulation", share=False, - max_worlds=64, + max_visible_envs=16, ) **Configuration options:** @@ -361,7 +371,6 @@ server, allowing you to view and interact with the scene from any browser. - ``share`` (bool, default ``False``): Request a public share URL from Viser for remote viewing. - ``record_to_viser`` (str or None, default ``None``): Path to save a ``.viser`` recording file. - ``verbose`` (bool, default ``True``): Print viewer server startup information. -- ``max_worlds`` (int or None, default ``None``): Maximum number of environments rendered. .. note:: @@ -371,7 +380,7 @@ server, allowing you to view and interact with the scene from any browser. Performance Note ---------------- -To reduce overhead when visualizing large-scale environments, consider: +When visualizing large-scale environments, consider: - Using Newton instead of Omniverse or Rerun - Reducing window sizes @@ -393,12 +402,6 @@ the num of environments can be overwritten and decreased using ``--num_envs``: python scripts/reinforcement_learning/rsl_rl/train.py --task Isaac-Cartpole-v0 --viz rerun --num_envs 512 -.. note:: - - A future feature will support visualizing only a subset of environments, which will improve visualization performance - and reduce resource usage while maintaining full-scale training in the background. - - **Rerun Visualizer FPS Control** The FPS control in the Rerun visualizer UI may not affect the visualization frame rate in all configurations. diff --git a/docs/source/how-to/record_video.rst b/docs/source/how-to/record_video.rst index aba743631295..01ee6240bb0c 100644 --- a/docs/source/how-to/record_video.rst +++ b/docs/source/how-to/record_video.rst @@ -3,6 +3,9 @@ Recording video clips during training Isaac Lab supports recording video clips during training using the `gymnasium.wrappers.RecordVideo `_ class. +When the ``--video`` flag is enabled, Isaac Lab captures a perspective view of the scene. The backend +is chosen automatically from the active physics and renderer stack: an Isaac Sim Kit camera or a +Newton GL headless viewer. This feature can be enabled by installing ``ffmpeg`` and using the following command line arguments with the training script: @@ -11,7 +14,6 @@ script: * ``--video_length``: length of each recorded video (in steps) * ``--video_interval``: interval between each video recording (in steps) -Make sure to also add the ``--enable_cameras`` argument when running headless. Note that enabling recording is equivalent to enabling rendering during training, which will slow down both startup and runtime performance. Example usage: @@ -23,3 +25,128 @@ Example usage: The recorded videos will be saved in the same directory as the training checkpoints, under ``IsaacLab/logs////videos/train``. + + +Overview +-------- + +The video recording feature is implemented using the ``VideoRecorder`` class. This class is responsible for resolving the video backend from the scene, capturing the video frames, and saving them to a file. + +* ``VideoRecorderCfg`` (``isaaclab.envs.utils.video_recorder_cfg``) holds resolution and world-space + perspective parameters ``camera_position`` and ``camera_target`` (defaults to a diagonal view of the + scene). +* ``VideoRecorder`` (``isaaclab.envs.utils.video_recorder``) picks a video backend from the scene + (Kit vs Newton GL), builds the matching low-level capture object, and returns RGB frames via + ``render_rgb_array()``. +* Direct RL, Direct MARL and manager-based RL environments copy the task's + :class:`~isaaclab.envs.common.ViewerCfg` ``eye`` and ``lookat`` into those fields before the + recorder is constructed, so training clips align with the task's intended viewport when + ``origin_type`` is ``"world"``. + + +Configuration: ``VideoRecorderCfg`` +------------------------------------ + +The dataclass lives in ``isaaclab.envs.utils.video_recorder_cfg``. Fields ``camera_position`` and +``camera_target`` are the perspective ``eye`` and ``lookat`` points in meters. + +.. literalinclude:: ../../../source/isaaclab/isaaclab/envs/utils/video_recorder_cfg.py + :language: python + :lines: 20-48 + + +Task framing: ``ViewerCfg`` +---------------------------- + +Tasks define the interactive viewer with :class:`~isaaclab.envs.common.ViewerCfg`. The ``eye`` and +``lookat`` tuples are the same values the RL base classes copy into ``VideoRecorderCfg`` (see below). +If your task uses ``origin_type="world"``, those tuples are world-space positions and match what the +perspective recorder expects. + +.. literalinclude:: ../../../source/isaaclab/isaaclab/envs/common.py + :language: python + :lines: 20-28 + + +Backend selection: Kit vs Newton GL +------------------------------------- + +``VideoRecorder`` resolves the implementation from the live :class:`~isaaclab.scene.InteractiveScene`. +If the user provides the PhysX physics (``presets=physx,...``) or Isaac RTX (``presets=isaac_rtx_renderer,...``) in the sensor stack, the Kit path is selected (``omni.replicator`` on +``/OmniverseKit_Persp``). The Newton GL path is selected when Newton physics is active (``presets=newton,...``) or the Newton +Warp renderer (``presets=newton_renderer,...``) appears in the sensor stack - and neither PhysX nor Isaac RTX is present to claim the +Kit path. OVRTX (``presets=ovrtx_renderer,...`` from ``isaaclab_ov``) can pair with IsaacSim or Newton physics; in that case the video backend is +selected via the physics preset. If both Kit and Newton GL signals are present (e.g., ``presets=physx,isaac_rtx_renderer,...`` or ``presets=newton,newton_renderer,...``), the Kit path is chosen. + +.. literalinclude:: ../../../source/isaaclab/isaaclab/envs/utils/video_recorder.py + :language: python + :lines: 38-59 + + +Construction and dispatch +-------------------------- + +When ``env_render_mode`` is ``"rgb_array"`` (as when wrappers or scripts request RGB frames for +video), the recorder instantiates the backend-specific helper and passes through ``camera_position``, +``camera_target``, and window size. + +.. literalinclude:: ../../../source/isaaclab/isaaclab/envs/utils/video_recorder.py + :language: python + :lines: 70-114 + + +Customising the camera view +---------------------------- + +When ``--video`` is passed, the recording camera uses the same +position and look-at target as the interactive viewer. The defaults come from +:class:`~isaaclab.envs.common.ViewerCfg`: + +* ``eye = (7.5, 7.5, 7.5)`` — camera position in world space (metres) +* ``lookat = (0.0, 0.0, 0.0)`` — camera look-at target in world space (metres) +* Resolution ``1280x720`` + +To change the recording angle, override the ``viewer`` field in your task's environment config. +The RL base classes automatically copy ``eye`` and ``lookat`` into ``VideoRecorderCfg`` before +recording starts (when ``origin_type`` is ``"world"``), so the video clip uses the same viewpoint +as the interactive viewport: + +.. code-block:: python + + from isaaclab.envs import ManagerBasedRLEnvCfg + from isaaclab.envs.common import ViewerCfg + from isaaclab.utils import configclass + + @configclass + class MyTaskCfg(ManagerBasedRLEnvCfg): + viewer: ViewerCfg = ViewerCfg( + eye=(5.0, 5.0, 5.0), + lookat=(0.0, 0.0, 1.0), + ) + + +Summary +------- + +.. list-table:: + :widths: 40 22 38 + :header-rows: 1 + + * - Stack example (``presets=...``) + - Video backend + - Capture mechanism + * - ``physx,...`` or ``isaac_rtx_renderer,...`` + - Kit (``"kit"``) + - ``/OmniverseKit_Persp`` + Replicator RGB + * - ``newton,...`` or ``newton_renderer,...`` (no Kit signals) + - Newton GL (``"newton_gl"``) + - ``newton.viewer.ViewerGL`` on the SDP Newton model + * - ``newton,...,ovrtx_renderer,...`` (OVRTX + Newton physics) + - Newton GL (``"newton_gl"``) + - ``newton.viewer.ViewerGL`` on the SDP Newton model + + +See also +-------- + +* :doc:`/source/features/visualization` - interactive visualizers diff --git a/docs/source/migration/migrating_to_isaaclab_3-0.rst b/docs/source/migration/migrating_to_isaaclab_3-0.rst index 12e0501c9028..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 ----------------- @@ -955,6 +998,14 @@ Common patterns that need updating: - ``isaaclab_physx`` * - :class:`~isaaclab_physx.sensors.FrameTransformer` - ``isaaclab_physx`` + * - :class:`~isaaclab.sensors.RayCaster` + - ``isaaclab`` + * - :class:`~isaaclab.sensors.RayCasterCamera` + - ``isaaclab`` + * - :class:`~isaaclab.sensors.MultiMeshRayCaster` + - ``isaaclab`` + * - :class:`~isaaclab.sensors.MultiMeshRayCasterCamera` + - ``isaaclab`` .. note:: @@ -974,6 +1025,129 @@ Common patterns that need updating: already passed to warp-native functions) should not be wrapped. +Ray Caster Warp Backend +~~~~~~~~~~~~~~~~~~~~~~~ + +The :class:`~isaaclab.sensors.RayCaster`, :class:`~isaaclab.sensors.RayCasterCamera`, +:class:`~isaaclab.sensors.MultiMeshRayCaster`, and +:class:`~isaaclab.sensors.MultiMeshRayCasterCamera` sensors have been transitioned from a +PyTorch/USD-based backend to a native Warp kernel pipeline. This improves performance by +eliminating per-step tensor allocations and torch-to-warp conversions, but introduces several +breaking changes. + + +RayCasterData Return Types +-------------------------- + +The :attr:`~isaaclab.sensors.RayCasterData.pos_w`, +:attr:`~isaaclab.sensors.RayCasterData.quat_w`, and +:attr:`~isaaclab.sensors.RayCasterData.ray_hits_w` properties now return ``wp.array`` instead of +``torch.Tensor``. This follows the same pattern as the general warp backend migration described +above. + +.. code-block:: python + + import warp as wp + + # Before (Isaac Lab 2.x) + ray_hits = ray_caster.data.ray_hits_w # torch.Tensor + sensor_pos = ray_caster.data.pos_w # torch.Tensor + + # After (Isaac Lab 3.x) + ray_hits = ray_caster.data.ray_hits_w # wp.array + sensor_pos = ray_caster.data.pos_w # wp.array + + # To use with torch operations, wrap with wp.to_torch() + ray_hits_torch = wp.to_torch(ray_caster.data.ray_hits_w) + sensor_pos_torch = wp.to_torch(ray_caster.data.pos_w) + + +Ray Alignment Configuration +---------------------------- + +The ``attach_yaw_only`` boolean parameter on :class:`~isaaclab.sensors.RayCasterCfg` has been +deprecated in favor of the new ``ray_alignment`` parameter, which accepts one of three string +values: + +.. list-table:: + :header-rows: 1 + :widths: 30 30 40 + + * - Old (2.x) + - New (3.0) + - Behavior + * - ``attach_yaw_only=False`` + - ``ray_alignment="base"`` + - Rays follow the full sensor orientation. + * - ``attach_yaw_only=True`` + - ``ray_alignment="yaw"`` + - Rays follow only the yaw component of the sensor orientation. + * - *(not available)* + - ``ray_alignment="world"`` + - Rays are always cast in the world frame (no rotation applied). + +.. code-block:: python + + # Before (Isaac Lab 2.x) + cfg = RayCasterCfg(attach_yaw_only=True, ...) + + # After (Isaac Lab 3.x) + cfg = RayCasterCfg(ray_alignment="yaw", ...) + + +Raycasting Kernel Signature Change +----------------------------------- + +The :func:`~isaaclab.utils.warp.kernels.raycast_dynamic_meshes_kernel` Warp kernel now requires +an ``env_mask`` parameter as its first argument. This is a ``wp.array(dtype=wp.bool)`` that +controls which environments are updated. The public Python wrapper +:func:`~isaaclab.utils.warp.ops.raycast_dynamic_meshes` has been updated to inject an all-True +mask automatically, so code using the wrapper is unaffected. + +If you call the kernel directly, update your launch call: + +.. code-block:: python + + import warp as wp + + # Before (Isaac Lab 2.x) + wp.launch( + raycast_dynamic_meshes_kernel, + dim=(num_meshes, num_envs, num_rays), + inputs=[ray_starts, ray_directions, mesh_ids, ...], + ) + + # After (Isaac Lab 3.x) -- env_mask is now the first input + env_mask = wp.ones(num_envs, dtype=wp.bool, device=device) + wp.launch( + raycast_dynamic_meshes_kernel, + dim=(num_meshes, num_envs, num_rays), + inputs=[env_mask, ray_starts, ray_directions, mesh_ids, ...], + ) + + +RayCaster.meshes Cache Key +-------------------------- + +The :attr:`~isaaclab.sensors.RayCaster.meshes` class variable, which caches warp meshes across +all :class:`~isaaclab.sensors.RayCaster` instances, is now keyed by ``(prim_path, device)`` tuples +instead of by ``prim_path`` alone. This prevents a mesh that was built on one device (e.g. CPU) +from being reused by a sensor running on a different device (e.g. CUDA), which caused illegal +memory accesses on systems without unified memory. + +Code that reads or writes this cache directly must update both the type annotation and the key: + +.. code-block:: python + + # Before (Isaac Lab 2.x) + meshes: ClassVar[dict[str, wp.Mesh]] = {} + wp_mesh = RayCaster.meshes[prim_path] + + # After (Isaac Lab 3.x) + meshes: ClassVar[dict[tuple[str, str], wp.Mesh]] = {} + wp_mesh = RayCaster.meshes[(prim_path, device)] + + Write Method Index/Mask Split ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ diff --git a/docs/source/overview/core-concepts/scene_data_providers.rst b/docs/source/overview/core-concepts/scene_data_providers.rst index 684dfcefcbef..527c59a03c12 100644 --- a/docs/source/overview/core-concepts/scene_data_providers.rst +++ b/docs/source/overview/core-concepts/scene_data_providers.rst @@ -27,9 +27,9 @@ The system has three layers: 1. **BaseSceneDataProvider** — abstract interface defining the contract: - - ``update(env_ids)`` — refresh cached scene data + - ``update()`` — refresh cached scene data (full Newton model/state sync when applicable) - ``get_newton_model()`` — return Newton model handle (if available) - - ``get_newton_state(env_ids)`` — return Newton state handle (if available) + - ``get_newton_state()`` — return Newton state handle (if available) - ``get_usd_stage()`` — return USD stage handle (if available) - ``get_transforms()`` — return body transforms - ``get_velocities()`` — return body velocities @@ -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/overview/environments.rst b/docs/source/overview/environments.rst index 95c29a33e8ea..80c12abb522e 100644 --- a/docs/source/overview/environments.rst +++ b/docs/source/overview/environments.rst @@ -43,54 +43,52 @@ Classic Classic environments that are based on IsaacGymEnvs implementation of MuJoCo-style environments. .. table:: - :widths: 33 37 30 - - +------------------+-----------------------------+-------------------------------------------------------------------------+ - | World | Environment ID | Description | - +==================+=============================+=========================================================================+ - | |humanoid| | |humanoid-link| | Move towards a direction with the MuJoCo humanoid robot | - | | | | - | | |humanoid-direct-link| | | - +------------------+-----------------------------+-------------------------------------------------------------------------+ - | |ant| | |ant-link| | Move towards a direction with the MuJoCo ant robot | - | | | | - | | |ant-direct-link| | | - +------------------+-----------------------------+-------------------------------------------------------------------------+ - | |cartpole| | |cartpole-link| | Move the cart to keep the pole upwards in the classic cartpole control | - | | | | - | | |cartpole-direct-link| | | - +------------------+-----------------------------+-------------------------------------------------------------------------+ - | |cartpole| | |cartpole-rgb-link| | Move the cart to keep the pole upwards in the classic cartpole control | - | | | and perceptive inputs. Requires running with ``--enable_cameras``. | - | | |cartpole-depth-link| | | - | | | | - | | |cartpole-rgb-direct-link| | | - | | | | - | | |cartpole-depth-direct-link|| | - +------------------+-----------------------------+-------------------------------------------------------------------------+ - | |cartpole| | |cartpole-resnet-link| | Move the cart to keep the pole upwards in the classic cartpole control | - | | | based off of features extracted from perceptive inputs with pre-trained | - | | |cartpole-theia-link| | frozen vision encoders. Requires running with ``--enable_cameras``. | - +------------------+-----------------------------+-------------------------------------------------------------------------+ + :widths: 25 30 25 20 + + +------------------+-----------------------------+-------------------------------------------------------------------------+-----------------------+ + | World | Environment ID | Description | Presets | + +==================+=============================+=========================================================================+=======================+ + | |humanoid| | |humanoid-link| | Move towards a direction with the MuJoCo humanoid robot | ``newton``, ``physx`` | + | | | | ``ovphysx`` | + | | |humanoid-direct-link| | | | + +------------------+-----------------------------+-------------------------------------------------------------------------+-----------------------+ + | |ant| | |ant-link| | Move towards a direction with the MuJoCo ant robot | ``newton``, ``physx`` | + | | | | ``ovphysx`` | + | | |ant-direct-link| | | | + +------------------+-----------------------------+-------------------------------------------------------------------------+-----------------------+ + | |cartpole| | |cartpole-link| | Move the cart to keep the pole upwards in the classic cartpole control | ``newton``, ``physx`` | + | | | | ``ovphysx`` | + | | |cartpole-direct-link| | | | + +------------------+-----------------------------+-------------------------------------------------------------------------+-----------------------+ + | |cartpole| | |cartpole-camera-presets| | Move the cart to keep the pole upwards in the classic cartpole control | ``newton``, ``physx`` | + | | | and perceptive inputs. Select data type via ``presets=``. Requires | ``newton_renderer``, | + | | | running with ``--enable_cameras``. | ``ovrtx_renderer``, | + | | | | ``rgb``, ``depth``, | + | | | | ``albedo``, | + | | | | ``semantic_`` | + | | | | ``segmentation``, | + | | | | ``simple_shading_*`` | + +------------------+-----------------------------+-------------------------------------------------------------------------+-----------------------+ + | |cartpole| | |cartpole-resnet-link| | Move the cart to keep the pole upwards in the classic cartpole control | ``newton``, ``physx`` | + | | | based off of features extracted from perceptive inputs with pre-trained | | + | | |cartpole-theia-link| | frozen vision encoders. Requires running with ``--enable_cameras``. | | + +------------------+-----------------------------+-------------------------------------------------------------------------+-----------------------+ .. |humanoid| image:: ../_static/tasks/classic/humanoid.jpg .. |ant| image:: ../_static/tasks/classic/ant.jpg .. |cartpole| image:: ../_static/tasks/classic/cartpole.jpg -.. |humanoid-link| replace:: `Isaac-Humanoid-v0 `__ -.. |ant-link| replace:: `Isaac-Ant-v0 `__ -.. |cartpole-link| replace:: `Isaac-Cartpole-v0 `__ -.. |cartpole-rgb-link| replace:: `Isaac-Cartpole-RGB-v0 `__ -.. |cartpole-depth-link| replace:: `Isaac-Cartpole-Depth-v0 `__ -.. |cartpole-resnet-link| replace:: `Isaac-Cartpole-RGB-ResNet18-v0 `__ -.. |cartpole-theia-link| replace:: `Isaac-Cartpole-RGB-TheiaTiny-v0 `__ +.. |humanoid-link| replace:: `Isaac-Humanoid-v0 <../../../source/isaaclab_tasks/isaaclab_tasks/manager_based/classic/humanoid/humanoid_env_cfg.py>`__ +.. |ant-link| replace:: `Isaac-Ant-v0 <../../../source/isaaclab_tasks/isaaclab_tasks/manager_based/classic/ant/ant_env_cfg.py>`__ +.. |cartpole-link| replace:: `Isaac-Cartpole-v0 <../../../source/isaaclab_tasks/isaaclab_tasks/manager_based/classic/cartpole/cartpole_env_cfg.py>`__ +.. |cartpole-camera-presets| replace:: `Isaac-Cartpole-Camera-Presets-Direct-v0 <../../../source/isaaclab_tasks/isaaclab_tasks/direct/cartpole/cartpole_camera_presets_env_cfg.py>`__ +.. |cartpole-resnet-link| replace:: `Isaac-Cartpole-RGB-ResNet18-v0 <../../../source/isaaclab_tasks/isaaclab_tasks/manager_based/classic/cartpole/cartpole_camera_env_cfg.py>`__ +.. |cartpole-theia-link| replace:: `Isaac-Cartpole-RGB-TheiaTiny-v0 <../../../source/isaaclab_tasks/isaaclab_tasks/manager_based/classic/cartpole/cartpole_camera_env_cfg.py>`__ -.. |humanoid-direct-link| replace:: `Isaac-Humanoid-Direct-v0 `__ -.. |ant-direct-link| replace:: `Isaac-Ant-Direct-v0 `__ -.. |cartpole-direct-link| replace:: `Isaac-Cartpole-Direct-v0 `__ -.. |cartpole-rgb-direct-link| replace:: `Isaac-Cartpole-RGB-Camera-Direct-v0 `__ -.. |cartpole-depth-direct-link| replace:: `Isaac-Cartpole-Depth-Camera-Direct-v0 `__ +.. |humanoid-direct-link| replace:: `Isaac-Humanoid-Direct-v0 <../../../source/isaaclab_tasks/isaaclab_tasks/direct/humanoid/humanoid_env.py>`__ +.. |ant-direct-link| replace:: `Isaac-Ant-Direct-v0 <../../../source/isaaclab_tasks/isaaclab_tasks/direct/ant/ant_env.py>`__ +.. |cartpole-direct-link| replace:: `Isaac-Cartpole-Direct-v0 <../../../source/isaaclab_tasks/isaaclab_tasks/direct/cartpole/cartpole_env.py>`__ Manipulation ~~~~~~~~~~~~ @@ -105,83 +103,107 @@ for the lift-cube environment: * |lift-cube-ik-rel-link|: Franka arm with relative IK control .. table:: - :widths: 33 37 30 - - +-------------------------+------------------------------+-----------------------------------------------------------------------------+ - | World | Environment ID | Description | - +=========================+==============================+=============================================================================+ - | |reach-franka| | |reach-franka-link| | Move the end-effector to a sampled target pose with the Franka robot | - +-------------------------+------------------------------+-----------------------------------------------------------------------------+ - | |reach-ur10| | |reach-ur10-link| | Move the end-effector to a sampled target pose with the UR10 robot | - +-------------------------+------------------------------+-----------------------------------------------------------------------------+ - | |deploy-reach-ur10e| | |deploy-reach-ur10e-link| | Move the end-effector to a sampled target pose with the UR10e robot | - | | | This policy has been deployed to a real robot | - +-------------------------+------------------------------+-----------------------------------------------------------------------------+ - | |lift-cube| | |lift-cube-link| | Pick a cube and bring it to a sampled target position with the Franka robot | - +-------------------------+------------------------------+-----------------------------------------------------------------------------+ - | |stack-cube| | |stack-cube-link| | Stack three cubes (bottom to top: blue, red, green) with the Franka robot. | - | | | Blueprint env used for the NVIDIA Isaac GR00T blueprint for synthetic | - | | |stack-cube-bp-link| | manipulation motion generation | - +-------------------------+------------------------------+-----------------------------------------------------------------------------+ - | |surface-gripper| | |long-suction-link| | Stack three cubes (bottom to top: blue, red, green) | - | | | with the UR10 arm and long surface gripper | - | | |short-suction-link| | or short surface gripper. | - +-------------------------+------------------------------+-----------------------------------------------------------------------------+ - | |cabi-franka| | |cabi-franka-link| | Grasp the handle of a cabinet's drawer and open it with the Franka robot | - | | | | - | | |franka-direct-link| | | - +-------------------------+------------------------------+-----------------------------------------------------------------------------+ - | |cube-allegro| | |cube-allegro-link| | In-hand reorientation of a cube using Allegro hand | - | | | | - | | |allegro-direct-link| | | - +-------------------------+------------------------------+-----------------------------------------------------------------------------+ - | |cube-shadow| | |cube-shadow-link| | In-hand reorientation of a cube using Shadow hand | - | | | | - | | |cube-shadow-ff-link| | | - | | | | - | | |cube-shadow-lstm-link| | | - +-------------------------+------------------------------+-----------------------------------------------------------------------------+ - | |cube-shadow| | |cube-shadow-vis-link| | In-hand reorientation of a cube using Shadow hand using perceptive inputs. | - | | | Requires running with ``--enable_cameras``. | - +-------------------------+------------------------------+-----------------------------------------------------------------------------+ - | |gr1_pick_place| | |gr1_pick_place-link| | Pick up and place an object in a basket with a GR-1 humanoid robot | - +-------------------------+------------------------------+-----------------------------------------------------------------------------+ - | |gr1_pp_waist| | |gr1_pp_waist-link| | Pick up and place an object in a basket with a GR-1 humanoid robot | - | | | with waist degrees-of-freedom enables that provides a wider reach space. | - +-------------------------+------------------------------+-----------------------------------------------------------------------------+ - | |g1_pick_place| | |g1_pick_place-link| | Pick up and place an object in a basket with a Unitree G1 humanoid robot | - +-------------------------+------------------------------+-----------------------------------------------------------------------------+ - | |g1_pick_place_fixed| | |g1_pick_place_fixed-link| | Pick up and place an object in a basket with a Unitree G1 humanoid robot | - | | | with three-fingered hands. Robot is set up with the base fixed in place. | - +-------------------------+------------------------------+-----------------------------------------------------------------------------+ - | |g1_pick_place_lm| | |g1_pick_place_lm-link| | Pick up and place an object in a basket with a Unitree G1 humanoid robot | - | | | with three-fingered hands and in-place locomanipulation capabilities | - | | | enabled (i.e. Robot lower body balances in-place while upper body is | - | | | controlled via Inverse Kinematics). | - +-------------------------+------------------------------+-----------------------------------------------------------------------------+ - | |kuka-allegro-lift| | |kuka-allegro-lift-link| | Pick up a primitive shape on the table and lift it to target position. | - | | | Supports state, single-camera, and dual-camera observation modes via | - | | | ``presets=single_camera`` / ``presets=duo_camera`` (see RL table below). | - +-------------------------+------------------------------+-----------------------------------------------------------------------------+ - | |kuka-allegro-reorient| | |kuka-allegro-reorient-link| | Pick up a primitive shape on the table and orient it to target pose. | - | | | Supports state, single-camera, and dual-camera observation modes via | - | | | ``presets=single_camera`` / ``presets=duo_camera`` (see RL table below). | - +-------------------------+------------------------------+-----------------------------------------------------------------------------+ - | |galbot_stack| | |galbot_stack-link| | Stack three cubes (bottom to top: blue, red, green) with the left arm of | - | | | a Galbot humanoid robot | - +-------------------------+------------------------------+-----------------------------------------------------------------------------+ - | |agibot_place_mug| | |agibot_place_mug-link| | Pick up and place a mug upright with a Agibot A2D humanoid robot | - +-------------------------+------------------------------+-----------------------------------------------------------------------------+ - | |agibot_place_toy| | |agibot_place_toy-link| | Pick up and place an object in a box with a Agibot A2D humanoid robot | - +-------------------------+------------------------------+-----------------------------------------------------------------------------+ - | |reach_openarm_bi| | |reach_openarm_bi-link| | Move the end-effector to sampled target poses with the OpenArm robot | - +-------------------------+------------------------------+-----------------------------------------------------------------------------+ - | |reach_openarm_uni| | |reach_openarm_uni-link| | Move the end-effector to a sampled target pose with the OpenArm robot | - +-------------------------+------------------------------+-----------------------------------------------------------------------------+ - | |lift_openarm_uni| | |lift_openarm_uni-link| | Pick a cube and bring it to a sampled target position with the OpenArm robot| - +-------------------------+------------------------------+-----------------------------------------------------------------------------+ - | |cabi_openarm_uni| | |cabi_openarm_uni-link| | Grasp the handle of a cabinet's drawer and open it with the OpenArm robot | - +-------------------------+------------------------------+-----------------------------------------------------------------------------+ + :widths: 25 30 25 20 + + +-------------------------+------------------------------+-----------------------------------------------------------------------------+-----------------------+ + | World | Environment ID | Description | Presets | + +=========================+==============================+=============================================================================+=======================+ + | |reach-franka| | |reach-franka-link| | Move the end-effector to a sampled target pose with the Franka robot | ``newton``, ``physx`` | + +-------------------------+------------------------------+-----------------------------------------------------------------------------+-----------------------+ + | |reach-ur10| | |reach-ur10-link| | Move the end-effector to a sampled target pose with the UR10 robot | ``newton``, ``physx`` | + +-------------------------+------------------------------+-----------------------------------------------------------------------------+-----------------------+ + | |deploy-reach-ur10e| | |deploy-reach-ur10e-link| | Move the end-effector to a sampled target pose with the UR10e robot | | + | | | This policy has been deployed to a real robot | | + +-------------------------+------------------------------+-----------------------------------------------------------------------------+-----------------------+ + | |lift-cube| | |lift-cube-link| | Pick a cube and bring it to a sampled target position with the Franka robot | | + +-------------------------+------------------------------+-----------------------------------------------------------------------------+-----------------------+ + | |stack-cube| | |stack-cube-link| | Stack three cubes (bottom to top: blue, red, green) with the Franka robot. | | + | | | Blueprint env used for the NVIDIA Isaac GR00T blueprint for synthetic | | + | | |stack-cube-bp-link| | manipulation motion generation | | + +-------------------------+------------------------------+-----------------------------------------------------------------------------+-----------------------+ + | |surface-gripper| | |long-suction-link| | Stack three cubes (bottom to top: blue, red, green) | | + | | | with the UR10 arm and long surface gripper | | + | | |short-suction-link| | or short surface gripper. | | + +-------------------------+------------------------------+-----------------------------------------------------------------------------+-----------------------+ + | |cabi-franka| | |cabi-franka-link| | Grasp the handle of a cabinet's drawer and open it with the Franka robot | ``newton``, ``physx`` | + | | | | | + | | |franka-direct-link| | | | + +-------------------------+------------------------------+-----------------------------------------------------------------------------+-----------------------+ + | |cube-allegro| | |cube-allegro-link| | In-hand reorientation of a cube using Allegro hand | ``newton``, ``physx`` | + | | | | | + | | |allegro-direct-link| | | | + +-------------------------+------------------------------+-----------------------------------------------------------------------------+-----------------------+ + | |cube-shadow| | |cube-shadow-link| | In-hand reorientation of a cube using Shadow hand | ``newton``, ``physx`` | + | | | | | + | | |cube-shadow-ff-link| | | | + | | | | | + | | |cube-shadow-lstm-link| | | | + +-------------------------+------------------------------+-----------------------------------------------------------------------------+-----------------------+ + | |cube-shadow| | |cube-shadow-vis-link| | In-hand reorientation of a cube using Shadow hand using perceptive inputs. | ``newton``, ``physx`` | + | | | Requires running with ``--enable_cameras``. | ``newton_renderer``, | + | | | | ``ovrtx_renderer``, | + | | | | ``rgb``, ``depth``, | + | | | | ``albedo``, ``full``, | + | | | | ``semantic_`` | + | | | | ``segmentation``, | + | | | | ``simple_shading_*`` | + +-------------------------+------------------------------+-----------------------------------------------------------------------------+-----------------------+ + | |gr1_pick_place| | |gr1_pick_place-link| | Pick up and place an object in a basket with a GR-1 humanoid robot | | + +-------------------------+------------------------------+-----------------------------------------------------------------------------+-----------------------+ + | |gr1_pp_waist| | |gr1_pp_waist-link| | Pick up and place an object in a basket with a GR-1 humanoid robot | | + | | | with waist degrees-of-freedom enables that provides a wider reach space. | | + +-------------------------+------------------------------+-----------------------------------------------------------------------------+-----------------------+ + | |g1_pick_place| | |g1_pick_place-link| | Pick up and place an object in a basket with a Unitree G1 humanoid robot | | + +-------------------------+------------------------------+-----------------------------------------------------------------------------+-----------------------+ + | |g1_pick_place_fixed| | |g1_pick_place_fixed-link| | Pick up and place an object in a basket with a Unitree G1 humanoid robot | | + | | | with three-fingered hands. Robot is set up with the base fixed in place. | | + +-------------------------+------------------------------+-----------------------------------------------------------------------------+-----------------------+ + | |g1_pick_place_lm| | |g1_pick_place_lm-link| | Pick up and place an object in a basket with a Unitree G1 humanoid robot | | + | | | with three-fingered hands and in-place locomanipulation capabilities | | + | | | enabled (i.e. Robot lower body balances in-place while upper body is | | + | | | controlled via Inverse Kinematics). | | + +-------------------------+------------------------------+-----------------------------------------------------------------------------+-----------------------+ + | |kuka-allegro-lift| | |kuka-allegro-lift-link| | Pick up a primitive shape on the table and lift it to target position. | ``newton``, ``physx`` | + | | | Supports state, single-camera, and dual-camera observation modes via | ``single_camera``, | + | | | ``presets=single_camera`` / ``presets=duo_camera`` (see RL table below). | ``duo_camera``, | + | | | | ``newton_renderer``, | + | | | | ``ovrtx_renderer``, | + | | | | ``rgb{64,128,256}``, | + | | | | ``depth{..}``, | + | | | | ``albedo{..}``, | + | | | | ``semantic_`` | + | | | | ``segmentation{..}``, | + | | | | ``simple_shading_*`` | + | | | | ``{64,128,256}`` | + +-------------------------+------------------------------+-----------------------------------------------------------------------------+-----------------------+ + | |kuka-allegro-reorient| | |kuka-allegro-reorient-link| | Pick up a primitive shape on the table and orient it to target pose. | ``newton``, ``physx`` | + | | | Supports state, single-camera, and dual-camera observation modes via | ``single_camera``, | + | | | ``presets=single_camera`` / ``presets=duo_camera`` (see RL table below). | ``duo_camera``, | + | | | | ``newton_renderer``, | + | | | | ``ovrtx_renderer``, | + | | | | ``rgb{64,128,256}``, | + | | | | ``depth{..}``, | + | | | | ``albedo{..}``, | + | | | | ``semantic_`` | + | | | | ``segmentation{..}``, | + | | | | ``simple_shading_*`` | + | | | | ``{64,128,256}`` | + +-------------------------+------------------------------+-----------------------------------------------------------------------------+-----------------------+ + | |galbot_stack| | |galbot_stack-link| | Stack three cubes (bottom to top: blue, red, green) with the left arm of | | + | | | a Galbot humanoid robot | | + +-------------------------+------------------------------+-----------------------------------------------------------------------------+-----------------------+ + | |agibot_place_mug| | |agibot_place_mug-link| | Pick up and place a mug upright with a Agibot A2D humanoid robot | | + +-------------------------+------------------------------+-----------------------------------------------------------------------------+-----------------------+ + | |agibot_place_toy| | |agibot_place_toy-link| | Pick up and place an object in a box with a Agibot A2D humanoid robot | | + +-------------------------+------------------------------+-----------------------------------------------------------------------------+-----------------------+ + | |reach_openarm_bi| | |reach_openarm_bi-link| | Move the end-effector to sampled target poses with the OpenArm robot | | + +-------------------------+------------------------------+-----------------------------------------------------------------------------+-----------------------+ + | |reach_openarm_uni| | |reach_openarm_uni-link| | Move the end-effector to a sampled target pose with the OpenArm robot | | + +-------------------------+------------------------------+-----------------------------------------------------------------------------+-----------------------+ + | |lift_openarm_uni| | |lift_openarm_uni-link| | Pick a cube and bring it to a sampled target position with the OpenArm robot| | + +-------------------------+------------------------------+-----------------------------------------------------------------------------+-----------------------+ + | |cabi_openarm_uni| | |cabi_openarm_uni-link| | Grasp the handle of a cabinet's drawer and open it with the OpenArm robot | | + +-------------------------+------------------------------+-----------------------------------------------------------------------------+-----------------------+ .. |reach-franka| image:: ../_static/tasks/manipulation/franka_reach.jpg .. |reach-ur10| image:: ../_static/tasks/manipulation/ur10_reach.jpg @@ -207,38 +229,38 @@ for the lift-cube environment: .. |lift_openarm_uni| image:: ../_static/tasks/manipulation/openarm_uni_lift.jpg .. |cabi_openarm_uni| image:: ../_static/tasks/manipulation/openarm_uni_open_drawer.jpg -.. |reach-franka-link| replace:: `Isaac-Reach-Franka-v0 `__ -.. |reach-ur10-link| replace:: `Isaac-Reach-UR10-v0 `__ -.. |deploy-reach-ur10e-link| replace:: `Isaac-Deploy-Reach-UR10e-v0 `__ -.. |lift-cube-link| replace:: `Isaac-Lift-Cube-Franka-v0 `__ -.. |lift-cube-ik-abs-link| replace:: `Isaac-Lift-Cube-Franka-IK-Abs-v0 `__ -.. |lift-cube-ik-rel-link| replace:: `Isaac-Lift-Cube-Franka-IK-Rel-v0 `__ -.. |cabi-franka-link| replace:: `Isaac-Open-Drawer-Franka-v0 `__ -.. |franka-direct-link| replace:: `Isaac-Franka-Cabinet-Direct-v0 `__ -.. |cube-allegro-link| replace:: `Isaac-Repose-Cube-Allegro-v0 `__ -.. |allegro-direct-link| replace:: `Isaac-Repose-Cube-Allegro-Direct-v0 `__ -.. |stack-cube-link| replace:: `Isaac-Stack-Cube-Franka-v0 `__ -.. |stack-cube-bp-link| replace:: `Isaac-Stack-Cube-Franka-IK-Rel-Blueprint-v0 `__ -.. |gr1_pick_place-link| replace:: `Isaac-PickPlace-GR1T2-Abs-v0 `__ -.. |g1_pick_place-link| replace:: `Isaac-PickPlace-G1-InspireFTP-Abs-v0 `__ -.. |g1_pick_place_fixed-link| replace:: `Isaac-PickPlace-FixedBaseUpperBodyIK-G1-Abs-v0 `__ -.. |g1_pick_place_lm-link| replace:: `Isaac-PickPlace-Locomanipulation-G1-Abs-v0 `__ -.. |long-suction-link| replace:: `Isaac-Stack-Cube-UR10-Long-Suction-IK-Rel-v0 `__ -.. |short-suction-link| replace:: `Isaac-Stack-Cube-UR10-Short-Suction-IK-Rel-v0 `__ -.. |gr1_pp_waist-link| replace:: `Isaac-PickPlace-GR1T2-WaistEnabled-Abs-v0 `__ -.. |galbot_stack-link| replace:: `Isaac-Stack-Cube-Galbot-Left-Arm-Gripper-RmpFlow-v0 `__ -.. |kuka-allegro-lift-link| replace:: `Isaac-Dexsuite-Kuka-Allegro-Lift-v0 `__ -.. |kuka-allegro-reorient-link| replace:: `Isaac-Dexsuite-Kuka-Allegro-Reorient-v0 `__ -.. |cube-shadow-link| replace:: `Isaac-Repose-Cube-Shadow-Direct-v0 `__ -.. |cube-shadow-ff-link| replace:: `Isaac-Repose-Cube-Shadow-OpenAI-FF-Direct-v0 `__ -.. |cube-shadow-lstm-link| replace:: `Isaac-Repose-Cube-Shadow-OpenAI-LSTM-Direct-v0 `__ -.. |cube-shadow-vis-link| replace:: `Isaac-Repose-Cube-Shadow-Vision-Direct-v0 `__ -.. |agibot_place_mug-link| replace:: `Isaac-Place-Mug-Agibot-Left-Arm-RmpFlow-v0 `__ -.. |agibot_place_toy-link| replace:: `Isaac-Place-Toy2Box-Agibot-Right-Arm-RmpFlow-v0 `__ -.. |reach_openarm_bi-link| replace:: `Isaac-Reach-OpenArm-Bi-v0 `__ -.. |reach_openarm_uni-link| replace:: `Isaac-Reach-OpenArm-v0 `__ -.. |lift_openarm_uni-link| replace:: `Isaac-Lift-Cube-OpenArm-v0 `__ -.. |cabi_openarm_uni-link| replace:: `Isaac-Open-Drawer-OpenArm-v0 `__ +.. |reach-franka-link| replace:: `Isaac-Reach-Franka-v0 <../../../source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/reach/config/franka/joint_pos_env_cfg.py>`__ +.. |reach-ur10-link| replace:: `Isaac-Reach-UR10-v0 <../../../source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/reach/config/ur_10/joint_pos_env_cfg.py>`__ +.. |deploy-reach-ur10e-link| replace:: `Isaac-Deploy-Reach-UR10e-v0 <../../../source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/deploy/reach/config/ur_10e/joint_pos_env_cfg.py>`__ +.. |lift-cube-link| replace:: `Isaac-Lift-Cube-Franka-v0 <../../../source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/lift/config/franka/joint_pos_env_cfg.py>`__ +.. |lift-cube-ik-abs-link| replace:: `Isaac-Lift-Cube-Franka-IK-Abs-v0 <../../../source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/lift/config/franka/ik_abs_env_cfg.py>`__ +.. |lift-cube-ik-rel-link| replace:: `Isaac-Lift-Cube-Franka-IK-Rel-v0 <../../../source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/lift/config/franka/ik_rel_env_cfg.py>`__ +.. |cabi-franka-link| replace:: `Isaac-Open-Drawer-Franka-v0 <../../../source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/cabinet/config/franka/joint_pos_env_cfg.py>`__ +.. |franka-direct-link| replace:: `Isaac-Franka-Cabinet-Direct-v0 <../../../source/isaaclab_tasks/isaaclab_tasks/direct/franka_cabinet/franka_cabinet_env.py>`__ +.. |cube-allegro-link| replace:: `Isaac-Repose-Cube-Allegro-v0 <../../../source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/inhand/config/allegro_hand/allegro_env_cfg.py>`__ +.. |allegro-direct-link| replace:: `Isaac-Repose-Cube-Allegro-Direct-v0 <../../../source/isaaclab_tasks/isaaclab_tasks/direct/allegro_hand/allegro_hand_env_cfg.py>`__ +.. |stack-cube-link| replace:: `Isaac-Stack-Cube-Franka-v0 <../../../source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/config/franka/stack_joint_pos_env_cfg.py>`__ +.. |stack-cube-bp-link| replace:: `Isaac-Stack-Cube-Franka-IK-Rel-Blueprint-v0 <../../../source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/config/franka/stack_ik_rel_blueprint_env_cfg.py>`__ +.. |gr1_pick_place-link| replace:: `Isaac-PickPlace-GR1T2-Abs-v0 <../../../source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/pick_place/pickplace_gr1t2_env_cfg.py>`__ +.. |g1_pick_place-link| replace:: `Isaac-PickPlace-G1-InspireFTP-Abs-v0 <../../../source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/pick_place/pickplace_unitree_g1_inspire_hand_env_cfg.py>`__ +.. |g1_pick_place_fixed-link| replace:: `Isaac-PickPlace-FixedBaseUpperBodyIK-G1-Abs-v0 <../../../source/isaaclab_tasks/isaaclab_tasks/manager_based/locomanipulation/pick_place/fixed_base_upper_body_ik_g1_env_cfg.py>`__ +.. |g1_pick_place_lm-link| replace:: `Isaac-PickPlace-Locomanipulation-G1-Abs-v0 <../../../source/isaaclab_tasks/isaaclab_tasks/manager_based/locomanipulation/pick_place/locomanipulation_g1_env_cfg.py>`__ +.. |long-suction-link| replace:: `Isaac-Stack-Cube-UR10-Long-Suction-IK-Rel-v0 <../../../source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/config/ur10_gripper/stack_ik_rel_env_cfg.py>`__ +.. |short-suction-link| replace:: `Isaac-Stack-Cube-UR10-Short-Suction-IK-Rel-v0 <../../../source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/config/ur10_gripper/stack_ik_rel_env_cfg.py>`__ +.. |gr1_pp_waist-link| replace:: `Isaac-PickPlace-GR1T2-WaistEnabled-Abs-v0 <../../../source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/pick_place/pickplace_gr1t2_waist_enabled_env_cfg.py>`__ +.. |galbot_stack-link| replace:: `Isaac-Stack-Cube-Galbot-Left-Arm-Gripper-RmpFlow-v0 <../../../source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/config/galbot/stack_rmp_rel_env_cfg.py>`__ +.. |kuka-allegro-lift-link| replace:: `Isaac-Dexsuite-Kuka-Allegro-Lift-v0 <../../../source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/config/kuka_allegro/dexsuite_kuka_allegro_env_cfg.py>`__ +.. |kuka-allegro-reorient-link| replace:: `Isaac-Dexsuite-Kuka-Allegro-Reorient-v0 <../../../source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/config/kuka_allegro/dexsuite_kuka_allegro_env_cfg.py>`__ +.. |cube-shadow-link| replace:: `Isaac-Repose-Cube-Shadow-Direct-v0 <../../../source/isaaclab_tasks/isaaclab_tasks/direct/shadow_hand/shadow_hand_env_cfg.py>`__ +.. |cube-shadow-ff-link| replace:: `Isaac-Repose-Cube-Shadow-OpenAI-FF-Direct-v0 <../../../source/isaaclab_tasks/isaaclab_tasks/direct/shadow_hand/shadow_hand_env_cfg.py>`__ +.. |cube-shadow-lstm-link| replace:: `Isaac-Repose-Cube-Shadow-OpenAI-LSTM-Direct-v0 <../../../source/isaaclab_tasks/isaaclab_tasks/direct/shadow_hand/shadow_hand_env_cfg.py>`__ +.. |cube-shadow-vis-link| replace:: `Isaac-Repose-Cube-Shadow-Vision-Direct-v0 <../../../source/isaaclab_tasks/isaaclab_tasks/direct/shadow_hand/shadow_hand_vision_env.py>`__ +.. |agibot_place_mug-link| replace:: `Isaac-Place-Mug-Agibot-Left-Arm-RmpFlow-v0 <../../../source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/place/config/agibot/place_upright_mug_rmp_rel_env_cfg.py>`__ +.. |agibot_place_toy-link| replace:: `Isaac-Place-Toy2Box-Agibot-Right-Arm-RmpFlow-v0 <../../../source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/place/config/agibot/place_toy2box_rmp_rel_env_cfg.py>`__ +.. |reach_openarm_bi-link| replace:: `Isaac-Reach-OpenArm-Bi-v0 <../../../source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/reach/config/openarm/bimanual/joint_pos_env_cfg.py>`__ +.. |reach_openarm_uni-link| replace:: `Isaac-Reach-OpenArm-v0 <../../../source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/reach/config/openarm/unimanual/joint_pos_env_cfg.py>`__ +.. |lift_openarm_uni-link| replace:: `Isaac-Lift-Cube-OpenArm-v0 <../../../source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/lift/config/openarm/joint_pos_env_cfg.py>`__ +.. |cabi_openarm_uni-link| replace:: `Isaac-Open-Drawer-OpenArm-v0 <../../../source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/cabinet/config/openarm/joint_pos_env_cfg.py>`__ Contact-rich Manipulation @@ -254,25 +276,25 @@ For example: * |factory-nut-link|: Nut-Bolt fastening with the Franka arm .. table:: - :widths: 33 37 30 - - +--------------------+-------------------------+-----------------------------------------------------------------------------+ - | World | Environment ID | Description | - +====================+=========================+=============================================================================+ - | |factory-peg| | |factory-peg-link| | Insert peg into the socket with the Franka robot | - +--------------------+-------------------------+-----------------------------------------------------------------------------+ - | |factory-gear| | |factory-gear-link| | Insert and mesh gear into the base with other gears, using the Franka robot | - +--------------------+-------------------------+-----------------------------------------------------------------------------+ - | |factory-nut| | |factory-nut-link| | Thread the nut onto the first 2 threads of the bolt, using the Franka robot | - +--------------------+-------------------------+-----------------------------------------------------------------------------+ + :widths: 25 30 25 20 + + +--------------------+-------------------------+-----------------------------------------------------------------------------+-----------------------+ + | World | Environment ID | Description | Presets | + +====================+=========================+=============================================================================+=======================+ + | |factory-peg| | |factory-peg-link| | Insert peg into the socket with the Franka robot | | + +--------------------+-------------------------+-----------------------------------------------------------------------------+-----------------------+ + | |factory-gear| | |factory-gear-link| | Insert and mesh gear into the base with other gears, using the Franka robot | | + +--------------------+-------------------------+-----------------------------------------------------------------------------+-----------------------+ + | |factory-nut| | |factory-nut-link| | Thread the nut onto the first 2 threads of the bolt, using the Franka robot | | + +--------------------+-------------------------+-----------------------------------------------------------------------------+-----------------------+ .. |factory-peg| image:: ../_static/tasks/factory/peg_insert.jpg .. |factory-gear| image:: ../_static/tasks/factory/gear_mesh.jpg .. |factory-nut| image:: ../_static/tasks/factory/nut_thread.jpg -.. |factory-peg-link| replace:: `Isaac-Factory-PegInsert-Direct-v0 `__ -.. |factory-gear-link| replace:: `Isaac-Factory-GearMesh-Direct-v0 `__ -.. |factory-nut-link| replace:: `Isaac-Factory-NutThread-Direct-v0 `__ +.. |factory-peg-link| replace:: `Isaac-Factory-PegInsert-Direct-v0 <../../../source/isaaclab_tasks/isaaclab_tasks/direct/factory/factory_env_cfg.py>`__ +.. |factory-gear-link| replace:: `Isaac-Factory-GearMesh-Direct-v0 <../../../source/isaaclab_tasks/isaaclab_tasks/direct/factory/factory_env_cfg.py>`__ +.. |factory-nut-link| replace:: `Isaac-Factory-NutThread-Direct-v0 <../../../source/isaaclab_tasks/isaaclab_tasks/direct/factory/factory_env_cfg.py>`__ AutoMate ~~~~~~~~ @@ -316,21 +338,21 @@ We provide environments for both disassembly and assembly. * To evaluate an assembly policy, we run the command ``python source/isaaclab_tasks/isaaclab_tasks/direct/automate/run_w_id.py --assembly_id=ASSEMBLY_ID --checkpoint=CHECKPOINT --log_eval``. The evaluation results are stored in ``evaluation_{ASSEMBLY_ID}.h5``. .. table:: - :widths: 33 37 30 + :widths: 25 30 25 20 - +--------------------+-------------------------+-----------------------------------------------------------------------------+ - | World | Environment ID | Description | - +====================+=========================+=============================================================================+ - | |disassembly| | |disassembly-link| | Lift a plug out of the socket with the Franka robot | - +--------------------+-------------------------+-----------------------------------------------------------------------------+ - | |assembly| | |assembly-link| | Insert a plug into its corresponding socket with the Franka robot | - +--------------------+-------------------------+-----------------------------------------------------------------------------+ + +--------------------+-------------------------+-----------------------------------------------------------------------------+-----------------------+ + | World | Environment ID | Description | Presets | + +====================+=========================+=============================================================================+=======================+ + | |disassembly| | |disassembly-link| | Lift a plug out of the socket with the Franka robot | | + +--------------------+-------------------------+-----------------------------------------------------------------------------+-----------------------+ + | |assembly| | |assembly-link| | Insert a plug into its corresponding socket with the Franka robot | | + +--------------------+-------------------------+-----------------------------------------------------------------------------+-----------------------+ .. |assembly| image:: ../_static/tasks/automate/00004.jpg .. |disassembly| image:: ../_static/tasks/automate/01053_disassembly.jpg -.. |assembly-link| replace:: `Isaac-AutoMate-Assembly-Direct-v0 `__ -.. |disassembly-link| replace:: `Isaac-AutoMate-Disassembly-Direct-v0 `__ +.. |assembly-link| replace:: `Isaac-AutoMate-Assembly-Direct-v0 <../../../source/isaaclab_tasks/isaaclab_tasks/direct/automate/assembly_env_cfg.py>`__ +.. |disassembly-link| replace:: `Isaac-AutoMate-Disassembly-Direct-v0 <../../../source/isaaclab_tasks/isaaclab_tasks/direct/automate/disassembly_env_cfg.py>`__ FORGE ~~~~~~~~ @@ -349,25 +371,25 @@ These tasks share the same task configurations and control options. You can swit * |forge-nut-link|: Nut-Bolt fastening with the Franka arm .. table:: - :widths: 33 37 30 - - +--------------------+-------------------------+-----------------------------------------------------------------------------+ - | World | Environment ID | Description | - +====================+=========================+=============================================================================+ - | |forge-peg| | |forge-peg-link| | Insert peg into the socket with the Franka robot | - +--------------------+-------------------------+-----------------------------------------------------------------------------+ - | |forge-gear| | |forge-gear-link| | Insert and mesh gear into the base with other gears, using the Franka robot | - +--------------------+-------------------------+-----------------------------------------------------------------------------+ - | |forge-nut| | |forge-nut-link| | Thread the nut onto the first 2 threads of the bolt, using the Franka robot | - +--------------------+-------------------------+-----------------------------------------------------------------------------+ + :widths: 25 30 25 20 + + +--------------------+-------------------------+-----------------------------------------------------------------------------+-----------------------+ + | World | Environment ID | Description | Presets | + +====================+=========================+=============================================================================+=======================+ + | |forge-peg| | |forge-peg-link| | Insert peg into the socket with the Franka robot | | + +--------------------+-------------------------+-----------------------------------------------------------------------------+-----------------------+ + | |forge-gear| | |forge-gear-link| | Insert and mesh gear into the base with other gears, using the Franka robot | | + +--------------------+-------------------------+-----------------------------------------------------------------------------+-----------------------+ + | |forge-nut| | |forge-nut-link| | Thread the nut onto the first 2 threads of the bolt, using the Franka robot | | + +--------------------+-------------------------+-----------------------------------------------------------------------------+-----------------------+ .. |forge-peg| image:: ../_static/tasks/factory/peg_insert.jpg .. |forge-gear| image:: ../_static/tasks/factory/gear_mesh.jpg .. |forge-nut| image:: ../_static/tasks/factory/nut_thread.jpg -.. |forge-peg-link| replace:: `Isaac-Forge-PegInsert-Direct-v0 `__ -.. |forge-gear-link| replace:: `Isaac-Forge-GearMesh-Direct-v0 `__ -.. |forge-nut-link| replace:: `Isaac-Forge-NutThread-Direct-v0 `__ +.. |forge-peg-link| replace:: `Isaac-Forge-PegInsert-Direct-v0 <../../../source/isaaclab_tasks/isaaclab_tasks/direct/forge/forge_env_cfg.py>`__ +.. |forge-gear-link| replace:: `Isaac-Forge-GearMesh-Direct-v0 <../../../source/isaaclab_tasks/isaaclab_tasks/direct/forge/forge_env_cfg.py>`__ +.. |forge-nut-link| replace:: `Isaac-Forge-NutThread-Direct-v0 <../../../source/isaaclab_tasks/isaaclab_tasks/direct/forge/forge_env_cfg.py>`__ Locomotion @@ -376,88 +398,88 @@ Locomotion Environments based on legged locomotion tasks. .. table:: - :widths: 33 37 30 - - +------------------------------+----------------------------------------------+------------------------------------------------------------------------------+ - | World | Environment ID | Description | - +==============================+==============================================+==============================================================================+ - | |velocity-flat-anymal-b| | |velocity-flat-anymal-b-link| | Track a velocity command on flat terrain with the Anymal B robot | - +------------------------------+----------------------------------------------+------------------------------------------------------------------------------+ - | |velocity-rough-anymal-b| | |velocity-rough-anymal-b-link| | Track a velocity command on rough terrain with the Anymal B robot | - +------------------------------+----------------------------------------------+------------------------------------------------------------------------------+ - | |velocity-flat-anymal-c| | |velocity-flat-anymal-c-link| | Track a velocity command on flat terrain with the Anymal C robot | - | | | | - | | |velocity-flat-anymal-c-direct-link| | | - +------------------------------+----------------------------------------------+------------------------------------------------------------------------------+ - | |velocity-rough-anymal-c| | |velocity-rough-anymal-c-link| | Track a velocity command on rough terrain with the Anymal C robot | - | | | | - | | |velocity-rough-anymal-c-direct-link| | | - +------------------------------+----------------------------------------------+------------------------------------------------------------------------------+ - | |velocity-flat-anymal-d| | |velocity-flat-anymal-d-link| | Track a velocity command on flat terrain with the Anymal D robot | - +------------------------------+----------------------------------------------+------------------------------------------------------------------------------+ - | |velocity-rough-anymal-d| | |velocity-rough-anymal-d-link| | Track a velocity command on rough terrain with the Anymal D robot | - +------------------------------+----------------------------------------------+------------------------------------------------------------------------------+ - | |velocity-flat-unitree-a1| | |velocity-flat-unitree-a1-link| | Track a velocity command on flat terrain with the Unitree A1 robot | - +------------------------------+----------------------------------------------+------------------------------------------------------------------------------+ - | |velocity-rough-unitree-a1| | |velocity-rough-unitree-a1-link| | Track a velocity command on rough terrain with the Unitree A1 robot | - +------------------------------+----------------------------------------------+------------------------------------------------------------------------------+ - | |velocity-flat-unitree-go1| | |velocity-flat-unitree-go1-link| | Track a velocity command on flat terrain with the Unitree Go1 robot | - +------------------------------+----------------------------------------------+------------------------------------------------------------------------------+ - | |velocity-rough-unitree-go1| | |velocity-rough-unitree-go1-link| | Track a velocity command on rough terrain with the Unitree Go1 robot | - +------------------------------+----------------------------------------------+------------------------------------------------------------------------------+ - | |velocity-flat-unitree-go2| | |velocity-flat-unitree-go2-link| | Track a velocity command on flat terrain with the Unitree Go2 robot | - +------------------------------+----------------------------------------------+------------------------------------------------------------------------------+ - | |velocity-rough-unitree-go2| | |velocity-rough-unitree-go2-link| | Track a velocity command on rough terrain with the Unitree Go2 robot | - +------------------------------+----------------------------------------------+------------------------------------------------------------------------------+ - | |velocity-flat-spot| | |velocity-flat-spot-link| | Track a velocity command on flat terrain with the Boston Dynamics Spot robot | - +------------------------------+----------------------------------------------+------------------------------------------------------------------------------+ - | |velocity-flat-h1| | |velocity-flat-h1-link| | Track a velocity command on flat terrain with the Unitree H1 robot | - +------------------------------+----------------------------------------------+------------------------------------------------------------------------------+ - | |velocity-rough-h1| | |velocity-rough-h1-link| | Track a velocity command on rough terrain with the Unitree H1 robot | - +------------------------------+----------------------------------------------+------------------------------------------------------------------------------+ - | |velocity-flat-g1| | |velocity-flat-g1-link| | Track a velocity command on flat terrain with the Unitree G1 robot | - +------------------------------+----------------------------------------------+------------------------------------------------------------------------------+ - | |velocity-rough-g1| | |velocity-rough-g1-link| | Track a velocity command on rough terrain with the Unitree G1 robot | - +------------------------------+----------------------------------------------+------------------------------------------------------------------------------+ - | |velocity-flat-digit| | |velocity-flat-digit-link| | Track a velocity command on flat terrain with the Agility Digit robot | - +------------------------------+----------------------------------------------+------------------------------------------------------------------------------+ - | |velocity-rough-digit| | |velocity-rough-digit-link| | Track a velocity command on rough terrain with the Agility Digit robot | - +------------------------------+----------------------------------------------+------------------------------------------------------------------------------+ - | |tracking-loco-manip-digit| | |tracking-loco-manip-digit-link| | Track a root velocity and hand pose command with the Agility Digit robot | - +------------------------------+----------------------------------------------+------------------------------------------------------------------------------+ - -.. |velocity-flat-anymal-b-link| replace:: `Isaac-Velocity-Flat-Anymal-B-v0 `__ -.. |velocity-rough-anymal-b-link| replace:: `Isaac-Velocity-Rough-Anymal-B-v0 `__ - -.. |velocity-flat-anymal-c-link| replace:: `Isaac-Velocity-Flat-Anymal-C-v0 `__ -.. |velocity-rough-anymal-c-link| replace:: `Isaac-Velocity-Rough-Anymal-C-v0 `__ - -.. |velocity-flat-anymal-c-direct-link| replace:: `Isaac-Velocity-Flat-Anymal-C-Direct-v0 `__ -.. |velocity-rough-anymal-c-direct-link| replace:: `Isaac-Velocity-Rough-Anymal-C-Direct-v0 `__ - -.. |velocity-flat-anymal-d-link| replace:: `Isaac-Velocity-Flat-Anymal-D-v0 `__ -.. |velocity-rough-anymal-d-link| replace:: `Isaac-Velocity-Rough-Anymal-D-v0 `__ - -.. |velocity-flat-unitree-a1-link| replace:: `Isaac-Velocity-Flat-Unitree-A1-v0 `__ -.. |velocity-rough-unitree-a1-link| replace:: `Isaac-Velocity-Rough-Unitree-A1-v0 `__ - -.. |velocity-flat-unitree-go1-link| replace:: `Isaac-Velocity-Flat-Unitree-Go1-v0 `__ -.. |velocity-rough-unitree-go1-link| replace:: `Isaac-Velocity-Rough-Unitree-Go1-v0 `__ - -.. |velocity-flat-unitree-go2-link| replace:: `Isaac-Velocity-Flat-Unitree-Go2-v0 `__ -.. |velocity-rough-unitree-go2-link| replace:: `Isaac-Velocity-Rough-Unitree-Go2-v0 `__ - -.. |velocity-flat-spot-link| replace:: `Isaac-Velocity-Flat-Spot-v0 `__ - -.. |velocity-flat-h1-link| replace:: `Isaac-Velocity-Flat-H1-v0 `__ -.. |velocity-rough-h1-link| replace:: `Isaac-Velocity-Rough-H1-v0 `__ - -.. |velocity-flat-g1-link| replace:: `Isaac-Velocity-Flat-G1-v0 `__ -.. |velocity-rough-g1-link| replace:: `Isaac-Velocity-Rough-G1-v0 `__ - -.. |velocity-flat-digit-link| replace:: `Isaac-Velocity-Flat-Digit-v0 `__ -.. |velocity-rough-digit-link| replace:: `Isaac-Velocity-Rough-Digit-v0 `__ -.. |tracking-loco-manip-digit-link| replace:: `Isaac-Tracking-LocoManip-Digit-v0 `__ + :widths: 25 30 25 20 + + +------------------------------+----------------------------------------------+------------------------------------------------------------------------------+-----------------------+ + | World | Environment ID | Description | Presets | + +==============================+==============================================+==============================================================================+=======================+ + | |velocity-flat-anymal-b| | |velocity-flat-anymal-b-link| | Track a velocity command on flat terrain with the Anymal B robot | ``newton``, ``physx`` | + +------------------------------+----------------------------------------------+------------------------------------------------------------------------------+-----------------------+ + | |velocity-rough-anymal-b| | |velocity-rough-anymal-b-link| | Track a velocity command on rough terrain with the Anymal B robot | ``newton``, ``physx`` | + +------------------------------+----------------------------------------------+------------------------------------------------------------------------------+-----------------------+ + | |velocity-flat-anymal-c| | |velocity-flat-anymal-c-link| | Track a velocity command on flat terrain with the Anymal C robot | ``newton``, ``physx`` | + | | | | | + | | |velocity-flat-anymal-c-direct-link| | | | + +------------------------------+----------------------------------------------+------------------------------------------------------------------------------+-----------------------+ + | |velocity-rough-anymal-c| | |velocity-rough-anymal-c-link| | Track a velocity command on rough terrain with the Anymal C robot | ``newton``, ``physx`` | + | | | | | + | | |velocity-rough-anymal-c-direct-link| | | | + +------------------------------+----------------------------------------------+------------------------------------------------------------------------------+-----------------------+ + | |velocity-flat-anymal-d| | |velocity-flat-anymal-d-link| | Track a velocity command on flat terrain with the Anymal D robot | ``newton``, ``physx`` | + +------------------------------+----------------------------------------------+------------------------------------------------------------------------------+-----------------------+ + | |velocity-rough-anymal-d| | |velocity-rough-anymal-d-link| | Track a velocity command on rough terrain with the Anymal D robot | ``newton``, ``physx`` | + +------------------------------+----------------------------------------------+------------------------------------------------------------------------------+-----------------------+ + | |velocity-flat-unitree-a1| | |velocity-flat-unitree-a1-link| | Track a velocity command on flat terrain with the Unitree A1 robot | ``newton``, ``physx`` | + +------------------------------+----------------------------------------------+------------------------------------------------------------------------------+-----------------------+ + | |velocity-rough-unitree-a1| | |velocity-rough-unitree-a1-link| | Track a velocity command on rough terrain with the Unitree A1 robot | ``newton``, ``physx`` | + +------------------------------+----------------------------------------------+------------------------------------------------------------------------------+-----------------------+ + | |velocity-flat-unitree-go1| | |velocity-flat-unitree-go1-link| | Track a velocity command on flat terrain with the Unitree Go1 robot | ``newton``, ``physx`` | + +------------------------------+----------------------------------------------+------------------------------------------------------------------------------+-----------------------+ + | |velocity-rough-unitree-go1| | |velocity-rough-unitree-go1-link| | Track a velocity command on rough terrain with the Unitree Go1 robot | ``newton``, ``physx`` | + +------------------------------+----------------------------------------------+------------------------------------------------------------------------------+-----------------------+ + | |velocity-flat-unitree-go2| | |velocity-flat-unitree-go2-link| | Track a velocity command on flat terrain with the Unitree Go2 robot | ``newton``, ``physx`` | + +------------------------------+----------------------------------------------+------------------------------------------------------------------------------+-----------------------+ + | |velocity-rough-unitree-go2| | |velocity-rough-unitree-go2-link| | Track a velocity command on rough terrain with the Unitree Go2 robot | ``newton``, ``physx`` | + +------------------------------+----------------------------------------------+------------------------------------------------------------------------------+-----------------------+ + | |velocity-flat-spot| | |velocity-flat-spot-link| | Track a velocity command on flat terrain with the Boston Dynamics Spot robot | ``newton``, ``physx`` | + +------------------------------+----------------------------------------------+------------------------------------------------------------------------------+-----------------------+ + | |velocity-flat-h1| | |velocity-flat-h1-link| | Track a velocity command on flat terrain with the Unitree H1 robot | ``newton``, ``physx`` | + +------------------------------+----------------------------------------------+------------------------------------------------------------------------------+-----------------------+ + | |velocity-rough-h1| | |velocity-rough-h1-link| | Track a velocity command on rough terrain with the Unitree H1 robot | ``newton``, ``physx`` | + +------------------------------+----------------------------------------------+------------------------------------------------------------------------------+-----------------------+ + | |velocity-flat-g1| | |velocity-flat-g1-link| | Track a velocity command on flat terrain with the Unitree G1 robot | ``newton``, ``physx`` | + +------------------------------+----------------------------------------------+------------------------------------------------------------------------------+-----------------------+ + | |velocity-rough-g1| | |velocity-rough-g1-link| | Track a velocity command on rough terrain with the Unitree G1 robot | ``newton``, ``physx`` | + +------------------------------+----------------------------------------------+------------------------------------------------------------------------------+-----------------------+ + | |velocity-flat-digit| | |velocity-flat-digit-link| | Track a velocity command on flat terrain with the Agility Digit robot | ``newton``, ``physx`` | + +------------------------------+----------------------------------------------+------------------------------------------------------------------------------+-----------------------+ + | |velocity-rough-digit| | |velocity-rough-digit-link| | Track a velocity command on rough terrain with the Agility Digit robot | ``newton``, ``physx`` | + +------------------------------+----------------------------------------------+------------------------------------------------------------------------------+-----------------------+ + | |tracking-loco-manip-digit| | |tracking-loco-manip-digit-link| | Track a root velocity and hand pose command with the Agility Digit robot | ``newton``, ``physx`` | + +------------------------------+----------------------------------------------+------------------------------------------------------------------------------+-----------------------+ + +.. |velocity-flat-anymal-b-link| replace:: `Isaac-Velocity-Flat-Anymal-B-v0 <../../../source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/config/anymal_b/flat_env_cfg.py>`__ +.. |velocity-rough-anymal-b-link| replace:: `Isaac-Velocity-Rough-Anymal-B-v0 <../../../source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/config/anymal_b/rough_env_cfg.py>`__ + +.. |velocity-flat-anymal-c-link| replace:: `Isaac-Velocity-Flat-Anymal-C-v0 <../../../source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/config/anymal_c/flat_env_cfg.py>`__ +.. |velocity-rough-anymal-c-link| replace:: `Isaac-Velocity-Rough-Anymal-C-v0 <../../../source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/config/anymal_c/rough_env_cfg.py>`__ + +.. |velocity-flat-anymal-c-direct-link| replace:: `Isaac-Velocity-Flat-Anymal-C-Direct-v0 <../../../source/isaaclab_tasks/isaaclab_tasks/direct/anymal_c/anymal_c_env.py>`__ +.. |velocity-rough-anymal-c-direct-link| replace:: `Isaac-Velocity-Rough-Anymal-C-Direct-v0 <../../../source/isaaclab_tasks/isaaclab_tasks/direct/anymal_c/anymal_c_env.py>`__ + +.. |velocity-flat-anymal-d-link| replace:: `Isaac-Velocity-Flat-Anymal-D-v0 <../../../source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/config/anymal_d/flat_env_cfg.py>`__ +.. |velocity-rough-anymal-d-link| replace:: `Isaac-Velocity-Rough-Anymal-D-v0 <../../../source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/config/anymal_d/rough_env_cfg.py>`__ + +.. |velocity-flat-unitree-a1-link| replace:: `Isaac-Velocity-Flat-Unitree-A1-v0 <../../../source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/config/a1/flat_env_cfg.py>`__ +.. |velocity-rough-unitree-a1-link| replace:: `Isaac-Velocity-Rough-Unitree-A1-v0 <../../../source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/config/a1/rough_env_cfg.py>`__ + +.. |velocity-flat-unitree-go1-link| replace:: `Isaac-Velocity-Flat-Unitree-Go1-v0 <../../../source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/config/go1/flat_env_cfg.py>`__ +.. |velocity-rough-unitree-go1-link| replace:: `Isaac-Velocity-Rough-Unitree-Go1-v0 <../../../source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/config/go1/rough_env_cfg.py>`__ + +.. |velocity-flat-unitree-go2-link| replace:: `Isaac-Velocity-Flat-Unitree-Go2-v0 <../../../source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/config/go2/flat_env_cfg.py>`__ +.. |velocity-rough-unitree-go2-link| replace:: `Isaac-Velocity-Rough-Unitree-Go2-v0 <../../../source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/config/go2/rough_env_cfg.py>`__ + +.. |velocity-flat-spot-link| replace:: `Isaac-Velocity-Flat-Spot-v0 <../../../source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/config/spot/flat_env_cfg.py>`__ + +.. |velocity-flat-h1-link| replace:: `Isaac-Velocity-Flat-H1-v0 <../../../source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/config/h1/flat_env_cfg.py>`__ +.. |velocity-rough-h1-link| replace:: `Isaac-Velocity-Rough-H1-v0 <../../../source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/config/h1/rough_env_cfg.py>`__ + +.. |velocity-flat-g1-link| replace:: `Isaac-Velocity-Flat-G1-v0 <../../../source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/config/g1/flat_env_cfg.py>`__ +.. |velocity-rough-g1-link| replace:: `Isaac-Velocity-Rough-G1-v0 <../../../source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/config/g1/rough_env_cfg.py>`__ + +.. |velocity-flat-digit-link| replace:: `Isaac-Velocity-Flat-Digit-v0 <../../../source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/config/digit/flat_env_cfg.py>`__ +.. |velocity-rough-digit-link| replace:: `Isaac-Velocity-Rough-Digit-v0 <../../../source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/config/digit/rough_env_cfg.py>`__ +.. |tracking-loco-manip-digit-link| replace:: `Isaac-Tracking-LocoManip-Digit-v0 <../../../source/isaaclab_tasks/isaaclab_tasks/manager_based/locomanipulation/tracking/config/digit/loco_manip_env_cfg.py>`__ .. |velocity-flat-anymal-b| image:: ../_static/tasks/locomotion/anymal_b_flat.jpg .. |velocity-rough-anymal-b| image:: ../_static/tasks/locomotion/anymal_b_rough.jpg @@ -484,15 +506,15 @@ Navigation ~~~~~~~~~~ .. table:: - :widths: 33 37 30 + :widths: 25 30 25 20 - +----------------+---------------------+-----------------------------------------------------------------------------+ - | World | Environment ID | Description | - +================+=====================+=============================================================================+ - | |anymal_c_nav| | |anymal_c_nav-link| | Navigate towards a target x-y position and heading with the ANYmal C robot. | - +----------------+---------------------+-----------------------------------------------------------------------------+ + +----------------+---------------------+-----------------------------------------------------------------------------+-----------------------+ + | World | Environment ID | Description | Presets | + +================+=====================+=============================================================================+=======================+ + | |anymal_c_nav| | |anymal_c_nav-link| | Navigate towards a target x-y position and heading with the ANYmal C robot. | ``newton``, ``physx`` | + +----------------+---------------------+-----------------------------------------------------------------------------+-----------------------+ -.. |anymal_c_nav-link| replace:: `Isaac-Navigation-Flat-Anymal-C-v0 `__ +.. |anymal_c_nav-link| replace:: `Isaac-Navigation-Flat-Anymal-C-v0 <../../../source/isaaclab_tasks/isaaclab_tasks/manager_based/navigation/config/anymal_c/navigation_env_cfg.py>`__ .. |anymal_c_nav| image:: ../_static/tasks/navigation/anymal_c_nav.jpg @@ -505,18 +527,18 @@ Multirotor See the `drone_arl` folder and the ARL robot config (`ARL_ROBOT_1_CFG`) in the codebase for details. -.. |arl_robot_track_position_state_based-link| replace:: `Isaac-TrackPositionNoObstacles-ARL-Robot-1-v0 `__ +.. |arl_robot_track_position_state_based-link| replace:: `Isaac-TrackPositionNoObstacles-ARL-Robot-1-v0 <../../../source/isaaclab_tasks/isaaclab_tasks/manager_based/drone_arl/track_position_state_based/config/arl_robot_1/track_position_state_based_env_cfg.py>`__ .. |arl_robot_track_position_state_based| image:: ../_static/tasks/drone_arl/arl_robot_1_track_position_state_based.jpg .. table:: - :widths: 33 37 30 + :widths: 25 30 25 20 - +----------------------------------------+---------------------------------------------+----------------------------------------------------------------------------------------+ - | World | Environment ID | Description | - +========================================+=============================================+========================================================================================+ - | |arl_robot_track_position_state_based| | |arl_robot_track_position_state_based-link| | Setpoint position control for the ARL robot using the track_position_state_based task. | - +----------------------------------------+---------------------------------------------+----------------------------------------------------------------------------------------+ + +----------------------------------------+---------------------------------------------+----------------------------------------------------------------------------------------+-----------------------+ + | World | Environment ID | Description | Presets | + +========================================+=============================================+========================================================================================+=======================+ + | |arl_robot_track_position_state_based| | |arl_robot_track_position_state_based-link| | Setpoint position control for the ARL robot using the track_position_state_based task. | | + +----------------------------------------+---------------------------------------------+----------------------------------------------------------------------------------------+-----------------------+ Others @@ -532,24 +554,24 @@ Others For evaluation, the play script's command line input ``--real-time`` allows the interaction loop between the environment and the agent to run in real time, if possible. .. table:: - :widths: 33 37 30 - - +----------------+---------------------------+-----------------------------------------------------------------------------+ - | World | Environment ID | Description | - +================+===========================+=============================================================================+ - | |quadcopter| | |quadcopter-link| | Fly and hover the Crazyflie copter at a goal point by applying thrust. | - +----------------+---------------------------+-----------------------------------------------------------------------------+ - | |humanoid_amp| | |humanoid_amp_dance-link| | Move a humanoid robot by imitating different pre-recorded human animations | - | | | (Adversarial Motion Priors). | - | | |humanoid_amp_run-link| | | - | | | | - | | |humanoid_amp_walk-link| | | - +----------------+---------------------------+-----------------------------------------------------------------------------+ - -.. |quadcopter-link| replace:: `Isaac-Quadcopter-Direct-v0 `__ -.. |humanoid_amp_dance-link| replace:: `Isaac-Humanoid-AMP-Dance-Direct-v0 `__ -.. |humanoid_amp_run-link| replace:: `Isaac-Humanoid-AMP-Run-Direct-v0 `__ -.. |humanoid_amp_walk-link| replace:: `Isaac-Humanoid-AMP-Walk-Direct-v0 `__ + :widths: 25 30 25 20 + + +----------------+---------------------------+-----------------------------------------------------------------------------+-----------------------+ + | World | Environment ID | Description | Presets | + +================+===========================+=============================================================================+=======================+ + | |quadcopter| | |quadcopter-link| | Fly and hover the Crazyflie copter at a goal point by applying thrust. | | + +----------------+---------------------------+-----------------------------------------------------------------------------+-----------------------+ + | |humanoid_amp| | |humanoid_amp_dance-link| | Move a humanoid robot by imitating different pre-recorded human animations | | + | | | (Adversarial Motion Priors). | | + | | |humanoid_amp_run-link| | | | + | | | | | + | | |humanoid_amp_walk-link| | | | + +----------------+---------------------------+-----------------------------------------------------------------------------+-----------------------+ + +.. |quadcopter-link| replace:: `Isaac-Quadcopter-Direct-v0 <../../../source/isaaclab_tasks/isaaclab_tasks/direct/quadcopter/quadcopter_env.py>`__ +.. |humanoid_amp_dance-link| replace:: `Isaac-Humanoid-AMP-Dance-Direct-v0 <../../../source/isaaclab_tasks/isaaclab_tasks/direct/humanoid_amp/humanoid_amp_env_cfg.py>`__ +.. |humanoid_amp_run-link| replace:: `Isaac-Humanoid-AMP-Run-Direct-v0 <../../../source/isaaclab_tasks/isaaclab_tasks/direct/humanoid_amp/humanoid_amp_env_cfg.py>`__ +.. |humanoid_amp_walk-link| replace:: `Isaac-Humanoid-AMP-Walk-Direct-v0 <../../../source/isaaclab_tasks/isaaclab_tasks/direct/humanoid_amp/humanoid_amp_env_cfg.py>`__ .. |quadcopter| image:: ../_static/tasks/others/quadcopter.jpg .. |humanoid_amp| image:: ../_static/tasks/others/humanoid_amp.jpg @@ -684,17 +706,17 @@ Classic ~~~~~~~ .. table:: - :widths: 33 37 30 + :widths: 25 30 25 20 - +------------------------+------------------------------------+-----------------------------------------------------------------------------------------------------------------------+ - | World | Environment ID | Description | - +========================+====================================+=======================================================================================================================+ - | |cart-double-pendulum| | |cart-double-pendulum-direct-link| | Move the cart and the pendulum to keep the last one upwards in the classic inverted double pendulum on a cart control | - +------------------------+------------------------------------+-----------------------------------------------------------------------------------------------------------------------+ + +------------------------+------------------------------------+-----------------------------------------------------------------------------------------------------------------------+-----------------------+ + | World | Environment ID | Description | Presets | + +========================+====================================+=======================================================================================================================+=======================+ + | |cart-double-pendulum| | |cart-double-pendulum-direct-link| | Move the cart and the pendulum to keep the last one upwards in the classic inverted double pendulum on a cart control | | + +------------------------+------------------------------------+-----------------------------------------------------------------------------------------------------------------------+-----------------------+ .. |cart-double-pendulum| image:: ../_static/tasks/classic/cart_double_pendulum.jpg -.. |cart-double-pendulum-direct-link| replace:: `Isaac-Cart-Double-Pendulum-Direct-v0 `__ +.. |cart-double-pendulum-direct-link| replace:: `Isaac-Cart-Double-Pendulum-Direct-v0 <../../../source/isaaclab_tasks/isaaclab_tasks/direct/cart_double_pendulum/cart_double_pendulum_env.py>`__ Manipulation ~~~~~~~~~~~~ @@ -702,17 +724,17 @@ Manipulation Environments based on fixed-arm manipulation tasks. .. table:: - :widths: 33 37 30 + :widths: 25 30 25 20 - +----------------------+--------------------------------+--------------------------------------------------------+ - | World | Environment ID | Description | - +======================+================================+========================================================+ - | |shadow-hand-over| | |shadow-hand-over-direct-link| | Passing an object from one hand over to the other hand | - +----------------------+--------------------------------+--------------------------------------------------------+ + +----------------------+--------------------------------+--------------------------------------------------------+-----------------------+ + | World | Environment ID | Description | Presets | + +======================+================================+========================================================+=======================+ + | |shadow-hand-over| | |shadow-hand-over-direct-link| | Passing an object from one hand over to the other hand | | + +----------------------+--------------------------------+--------------------------------------------------------+-----------------------+ .. |shadow-hand-over| image:: ../_static/tasks/manipulation/shadow_hand_over.jpg -.. |shadow-hand-over-direct-link| replace:: `Isaac-Shadow-Hand-Over-Direct-v0 `__ +.. |shadow-hand-over-direct-link| replace:: `Isaac-Shadow-Hand-Over-Direct-v0 <../../../source/isaaclab_tasks/isaaclab_tasks/direct/shadow_hand_over/shadow_hand_over_env.py>`__ | @@ -724,308 +746,368 @@ provided when running ``play.py`` or any inferencing workflows. These tasks prov inferencing, including reading from an already trained checkpoint and disabling runtime perturbations used for training. .. list-table:: - :widths: 33 25 19 25 + :widths: 28 20 13 22 17 * - **Task Name** - **Inference Task Name** - **Workflow** - **RL Library** + - **Presets** * - Isaac-Ant-Direct-v0 - - Direct - **rl_games** (PPO), **rsl_rl** (PPO), **skrl** (PPO) + - ``newton``, ``physx``, ``ovphysx`` * - Isaac-Ant-v0 - - Manager Based - **rsl_rl** (PPO), **rl_games** (PPO), **skrl** (PPO), **sb3** (PPO) + - ``newton``, ``physx`` * - Isaac-Cart-Double-Pendulum-Direct-v0 - - Direct - **rl_games** (PPO), **skrl** (IPPO, PPO, MAPPO) + - * - Isaac-Cartpole-Camera-Showcase-Box-Box-Direct-v0 (Requires running with ``--enable_cameras``) - - Direct - **skrl** (PPO) + - ``newton_renderer``, ``ovrtx_renderer``, ``isaacsim_rtx_renderer`` * - Isaac-Cartpole-Camera-Showcase-Box-Discrete-Direct-v0 (Requires running with ``--enable_cameras``) - - Direct - **skrl** (PPO) + - ``newton_renderer``, ``ovrtx_renderer``, ``isaacsim_rtx_renderer`` * - Isaac-Cartpole-Camera-Showcase-Box-MultiDiscrete-Direct-v0 (Requires running with ``--enable_cameras``) - - Direct - **skrl** (PPO) + - ``newton_renderer``, ``ovrtx_renderer``, ``isaacsim_rtx_renderer`` * - Isaac-Cartpole-Camera-Showcase-Dict-Box-Direct-v0 (Requires running with ``--enable_cameras``) - - Direct - **skrl** (PPO) + - ``newton_renderer``, ``ovrtx_renderer``, ``isaacsim_rtx_renderer`` * - Isaac-Cartpole-Camera-Showcase-Dict-Discrete-Direct-v0 (Requires running with ``--enable_cameras``) - - Direct - **skrl** (PPO) + - ``newton_renderer``, ``ovrtx_renderer``, ``isaacsim_rtx_renderer`` * - Isaac-Cartpole-Camera-Showcase-Dict-MultiDiscrete-Direct-v0 (Requires running with ``--enable_cameras``) - - Direct - **skrl** (PPO) + - ``newton_renderer``, ``ovrtx_renderer``, ``isaacsim_rtx_renderer`` * - Isaac-Cartpole-Camera-Showcase-Tuple-Box-Direct-v0 (Requires running with ``--enable_cameras``) - - Direct - **skrl** (PPO) + - ``newton_renderer``, ``ovrtx_renderer``, ``isaacsim_rtx_renderer`` * - Isaac-Cartpole-Camera-Showcase-Tuple-Discrete-Direct-v0 (Requires running with ``--enable_cameras``) - - Direct - **skrl** (PPO) + - ``newton_renderer``, ``ovrtx_renderer``, ``isaacsim_rtx_renderer`` * - Isaac-Cartpole-Camera-Showcase-Tuple-MultiDiscrete-Direct-v0 (Requires running with ``--enable_cameras``) - - Direct - **skrl** (PPO) - * - Isaac-Cartpole-Depth-Camera-Direct-v0 (Requires running with ``--enable_cameras``) + - ``newton_renderer``, ``ovrtx_renderer``, ``isaacsim_rtx_renderer`` + * - Isaac-Cartpole-Camera-Presets-Direct-v0 (Requires running with ``--enable_cameras``) - - Direct - **rl_games** (PPO), **skrl** (PPO) - * - Isaac-Cartpole-Depth-v0 (Requires running with ``--enable_cameras``) - - - - Manager Based - - **rl_games** (PPO) + - ``newton``, ``physx``, ``newton_renderer``, ``ovrtx_renderer``, ``isaacsim_rtx_renderer``, ``rgb``, ``depth``, ``albedo``, ``semantic_segmentation``, ``simple_shading_constant_diffuse``, ``simple_shading_diffuse_mdl``, ``simple_shading_full_mdl`` * - Isaac-Cartpole-Direct-v0 - - Direct - **rl_games** (PPO), **rsl_rl** (PPO), **skrl** (PPO), **sb3** (PPO) - * - Isaac-Cartpole-RGB-Camera-Direct-v0 (Requires running with ``--enable_cameras``) - - - - Direct - - **rl_games** (PPO), **skrl** (PPO) + - ``newton``, ``physx``, ``ovphysx`` * - Isaac-Cartpole-RGB-ResNet18-v0 (Requires running with ``--enable_cameras``) - - Manager Based - **rl_games** (PPO) + - ``newton``, ``physx`` * - Isaac-Cartpole-RGB-TheiaTiny-v0 (Requires running with ``--enable_cameras``) - - Manager Based - **rl_games** (PPO) - * - Isaac-Cartpole-RGB-v0 (Requires running with ``--enable_cameras``) - - - - Manager Based - - **rl_games** (PPO) + - ``newton``, ``physx`` * - Isaac-Cartpole-Showcase-Box-Box-Direct-v0 - - Direct - **skrl** (PPO) + - ``newton``, ``physx``, ``ovphysx`` * - Isaac-Cartpole-Showcase-Box-Discrete-Direct-v0 - - Direct - **skrl** (PPO) + - ``newton``, ``physx``, ``ovphysx`` * - Isaac-Cartpole-Showcase-Box-MultiDiscrete-Direct-v0 - - Direct - **skrl** (PPO) + - ``newton``, ``physx``, ``ovphysx`` * - Isaac-Cartpole-Showcase-Dict-Box-Direct-v0 - - Direct - **skrl** (PPO) + - ``newton``, ``physx``, ``ovphysx`` * - Isaac-Cartpole-Showcase-Dict-Discrete-Direct-v0 - - Direct - **skrl** (PPO) + - ``newton``, ``physx``, ``ovphysx`` * - Isaac-Cartpole-Showcase-Dict-MultiDiscrete-Direct-v0 - - Direct - **skrl** (PPO) + - ``newton``, ``physx``, ``ovphysx`` * - Isaac-Cartpole-Showcase-Discrete-Box-Direct-v0 - - Direct - **skrl** (PPO) + - ``newton``, ``physx``, ``ovphysx`` * - Isaac-Cartpole-Showcase-Discrete-Discrete-Direct-v0 - - Direct - **skrl** (PPO) + - ``newton``, ``physx``, ``ovphysx`` * - Isaac-Cartpole-Showcase-Discrete-MultiDiscrete-Direct-v0 - - Direct - **skrl** (PPO) + - ``newton``, ``physx``, ``ovphysx`` * - Isaac-Cartpole-Showcase-MultiDiscrete-Box-Direct-v0 - - Direct - **skrl** (PPO) + - ``newton``, ``physx``, ``ovphysx`` * - Isaac-Cartpole-Showcase-MultiDiscrete-Discrete-Direct-v0 - - Direct - **skrl** (PPO) + - ``newton``, ``physx``, ``ovphysx`` * - Isaac-Cartpole-Showcase-MultiDiscrete-MultiDiscrete-Direct-v0 - - Direct - **skrl** (PPO) + - ``newton``, ``physx``, ``ovphysx`` * - Isaac-Cartpole-Showcase-Tuple-Box-Direct-v0 - - Direct - **skrl** (PPO) + - ``newton``, ``physx``, ``ovphysx`` * - Isaac-Cartpole-Showcase-Tuple-Discrete-Direct-v0 - - Direct - **skrl** (PPO) + - ``newton``, ``physx``, ``ovphysx`` * - Isaac-Cartpole-Showcase-Tuple-MultiDiscrete-Direct-v0 - - Direct - **skrl** (PPO) + - ``newton``, ``physx``, ``ovphysx`` * - Isaac-Cartpole-v0 - - Manager Based - **rl_games** (PPO), **rsl_rl** (PPO), **skrl** (PPO), **sb3** (PPO) + - ``newton``, ``physx`` * - Isaac-Factory-GearMesh-Direct-v0 - - Direct - **rl_games** (PPO) + - * - Isaac-Factory-NutThread-Direct-v0 - - Direct - **rl_games** (PPO) + - * - Isaac-Factory-PegInsert-Direct-v0 - - Direct - **rl_games** (PPO) + - * - Isaac-AutoMate-Assembly-Direct-v0 - - Direct - **rl_games** (PPO) + - * - Isaac-AutoMate-Disassembly-Direct-v0 - - Direct - + - * - Isaac-Forge-GearMesh-Direct-v0 - - Direct - **rl_games** (PPO) + - * - Isaac-Forge-NutThread-Direct-v0 - - Direct - **rl_games** (PPO) + - * - Isaac-Forge-PegInsert-Direct-v0 - - Direct - **rl_games** (PPO) + - * - Isaac-Franka-Cabinet-Direct-v0 - - Direct - **rl_games** (PPO), **rsl_rl** (PPO), **skrl** (PPO) + - * - Isaac-Humanoid-AMP-Dance-Direct-v0 - - Direct - **skrl** (AMP) + - * - Isaac-Humanoid-AMP-Run-Direct-v0 - - Direct - **skrl** (AMP) + - * - Isaac-Humanoid-AMP-Walk-Direct-v0 - - Direct - **skrl** (AMP) + - * - Isaac-Humanoid-Direct-v0 - - Direct - **rl_games** (PPO), **rsl_rl** (PPO), **skrl** (PPO) + - ``newton``, ``physx``, ``ovphysx`` * - Isaac-Humanoid-v0 - - Manager Based - **rsl_rl** (PPO), **rl_games** (PPO), **skrl** (PPO), **sb3** (PPO) + - ``newton``, ``physx`` * - Isaac-Lift-Cube-Franka-IK-Abs-v0 - - Manager Based - + - * - Isaac-Lift-Cube-Franka-IK-Rel-v0 - - Manager Based - + - * - Isaac-Lift-Cube-Franka-v0 - Isaac-Lift-Cube-Franka-Play-v0 - Manager Based - **rsl_rl** (PPO), **skrl** (PPO), **rl_games** (PPO), **sb3** (PPO) + - * - Isaac-Lift-Teddy-Bear-Franka-IK-Abs-v0 - - Manager Based - + - * - Isaac-Tracking-LocoManip-Digit-v0 - Isaac-Tracking-LocoManip-Digit-Play-v0 - Manager Based - **rsl_rl** (PPO) + - ``newton``, ``physx`` * - Isaac-Navigation-Flat-Anymal-C-v0 - Isaac-Navigation-Flat-Anymal-C-Play-v0 - Manager Based - **rsl_rl** (PPO), **skrl** (PPO) + - ``newton``, ``physx`` * - Isaac-Open-Drawer-Franka-IK-Abs-v0 - - Manager Based - + - * - Isaac-Open-Drawer-Franka-IK-Rel-v0 - - Manager Based - + - * - Isaac-Open-Drawer-Franka-v0 - Isaac-Open-Drawer-Franka-Play-v0 - Manager Based - **rsl_rl** (PPO), **rl_games** (PPO), **skrl** (PPO) + - ``newton``, ``physx`` * - Isaac-Quadcopter-Direct-v0 - - Direct - **rl_games** (PPO), **rsl_rl** (PPO), **skrl** (PPO) + - * - Isaac-Reach-Franka-IK-Abs-v0 - - Manager Based - + - * - Isaac-Reach-Franka-IK-Rel-v0 - - Manager Based - + - * - Isaac-Reach-Franka-OSC-v0 - Isaac-Reach-Franka-OSC-Play-v0 - Manager Based - **rsl_rl** (PPO) + - ``newton``, ``physx`` * - Isaac-Reach-Franka-v0 - Isaac-Reach-Franka-Play-v0 - Manager Based - **rl_games** (PPO), **rsl_rl** (PPO), **skrl** (PPO) + - ``newton``, ``physx`` * - Isaac-Reach-UR10-v0 - Isaac-Reach-UR10-Play-v0 - Manager Based - **rl_games** (PPO), **rsl_rl** (PPO), **skrl** (PPO) + - ``newton``, ``physx`` * - Isaac-Deploy-Reach-UR10e-v0 - Isaac-Deploy-Reach-UR10e-Play-v0 - Manager Based - **rsl_rl** (PPO) + - * - Isaac-Repose-Cube-Allegro-Direct-v0 - - Direct - **rl_games** (PPO), **rsl_rl** (PPO), **skrl** (PPO) + - ``newton``, ``physx`` * - Isaac-Repose-Cube-Allegro-NoVelObs-v0 - Isaac-Repose-Cube-Allegro-NoVelObs-Play-v0 - Manager Based - **rsl_rl** (PPO), **rl_games** (PPO), **skrl** (PPO) + - * - Isaac-Repose-Cube-Allegro-v0 - Isaac-Repose-Cube-Allegro-Play-v0 - Manager Based - **rsl_rl** (PPO), **rl_games** (PPO), **skrl** (PPO) + - * - Isaac-Repose-Cube-Shadow-Direct-v0 - - Direct - **rl_games** (PPO), **rsl_rl** (PPO), **skrl** (PPO) + - ``newton``, ``physx`` * - Isaac-Repose-Cube-Shadow-OpenAI-FF-Direct-v0 - - Direct - **rl_games** (FF), **rsl_rl** (PPO), **skrl** (PPO) + - ``newton``, ``physx`` * - Isaac-Repose-Cube-Shadow-OpenAI-LSTM-Direct-v0 - - Direct - **rl_games** (LSTM) + - ``newton``, ``physx`` * - Isaac-Repose-Cube-Shadow-Vision-Direct-v0 (Requires running with ``--enable_cameras``) - Isaac-Repose-Cube-Shadow-Vision-Direct-Play-v0 (Requires running with ``--enable_cameras``) - Direct - **rsl_rl** (PPO), **rl_games** (VISION) + - ``newton``, ``physx``, ``newton_renderer``, ``ovrtx_renderer``, ``isaacsim_rtx_renderer``, ``rgb``, ``depth``, ``albedo``, ``full``, ``semantic_segmentation``, ``simple_shading_constant_diffuse``, ``simple_shading_diffuse_mdl``, ``simple_shading_full_mdl`` * - Isaac-Shadow-Hand-Over-Direct-v0 - - Direct - **rl_games** (PPO), **skrl** (IPPO, PPO, MAPPO) + - * - Isaac-Stack-Cube-Franka-IK-Rel-v0 - - Manager Based - + - * - Isaac-Dexsuite-Kuka-Allegro-Lift-v0 Camera variants (requires ``--enable_cameras``): @@ -1041,6 +1123,7 @@ inferencing, including reading from an already trained checkpoint and disabling - Isaac-Dexsuite-Kuka-Allegro-Lift-Play-v0 - Manager Based - **rl_games** (PPO), **rsl_rl** (PPO) + - ``newton``, ``physx``, ``single_camera``, ``duo_camera``, ``state``, ``newton_renderer``, ``ovrtx_renderer``, ``isaacsim_rtx_renderer``, ``rgb64``, ``rgb128``, ``rgb256``, ``depth64``, ``depth128``, ``depth256``, ``albedo64``, ``albedo128``, ``albedo256``, ``semantic_segmentation64``, ``semantic_segmentation128``, ``semantic_segmentation256``, ``simple_shading_constant_diffuse64``, ``simple_shading_constant_diffuse128``, ``simple_shading_constant_diffuse256``, ``simple_shading_diffuse_mdl64``, ``simple_shading_diffuse_mdl128``, ``simple_shading_diffuse_mdl256``, ``simple_shading_full_mdl64``, ``simple_shading_full_mdl128``, ``simple_shading_full_mdl256`` * - Isaac-Dexsuite-Kuka-Allegro-Reorient-v0 Camera variants (requires ``--enable_cameras``): @@ -1053,176 +1136,220 @@ inferencing, including reading from an already trained checkpoint and disabling - Isaac-Dexsuite-Kuka-Allegro-Reorient-Play-v0 - Manager Based - **rl_games** (PPO), **rsl_rl** (PPO) + - ``newton``, ``physx``, ``single_camera``, ``duo_camera``, ``state``, ``newton_renderer``, ``ovrtx_renderer``, ``isaacsim_rtx_renderer``, ``rgb64``, ``rgb128``, ``rgb256``, ``depth64``, ``depth128``, ``depth256``, ``albedo64``, ``albedo128``, ``albedo256``, ``semantic_segmentation64``, ``semantic_segmentation128``, ``semantic_segmentation256``, ``simple_shading_constant_diffuse64``, ``simple_shading_constant_diffuse128``, ``simple_shading_constant_diffuse256``, ``simple_shading_diffuse_mdl64``, ``simple_shading_diffuse_mdl128``, ``simple_shading_diffuse_mdl256``, ``simple_shading_full_mdl64``, ``simple_shading_full_mdl128``, ``simple_shading_full_mdl256`` * - Isaac-Stack-Cube-Franka-v0 - - Manager Based - + - * - Isaac-Stack-Cube-Instance-Randomize-Franka-IK-Rel-v0 - - Manager Based - + - * - Isaac-Stack-Cube-Instance-Randomize-Franka-v0 - - Manager Based - + - * - Isaac-PickPlace-G1-InspireFTP-Abs-v0 - - Manager Based - + - * - Isaac-Stack-Cube-UR10-Long-Suction-IK-Rel-v0 - - Manager Based - + - * - Isaac-Stack-Cube-UR10-Short-Suction-IK-Rel-v0 - - Manager Based - + - * - Isaac-Stack-Cube-Galbot-Left-Arm-Gripper-RmpFlow-v0 - - Manager Based - + - * - Isaac-Stack-Cube-Galbot-Right-Arm-Suction-RmpFlow-v0 - - Manager Based - + - * - Isaac-Stack-Cube-Galbot-Left-Arm-Gripper-Visuomotor-v0 - Isaac-Stack-Cube-Galbot-Left-Arm-Gripper-Visuomotor-Play-v0 - Manager Based - + - * - Isaac-Place-Mug-Agibot-Left-Arm-RmpFlow-v0 - - Manager Based - + - * - Isaac-Place-Toy2Box-Agibot-Right-Arm-RmpFlow-v0 - - Manager Based - + - * - Isaac-Stack-Cube-Galbot-Left-Arm-Gripper-RmpFlow-v0 - - Manager Based - + - * - Isaac-Stack-Cube-Galbot-Right-Arm-Suction-RmpFlow-v0 - - Manager Based - + - * - Isaac-Stack-Cube-Galbot-Left-Arm-Gripper-Visuomotor-v0 - Isaac-Stack-Cube-Galbot-Left-Arm-Gripper-Visuomotor-Play-v0 - Manager Based - + - * - Isaac-Place-Mug-Agibot-Left-Arm-RmpFlow-v0 - - Manager Based - + - * - Isaac-Place-Toy2Box-Agibot-Right-Arm-RmpFlow-v0 - - Manager Based - + - * - Isaac-Velocity-Flat-Anymal-B-v0 - Isaac-Velocity-Flat-Anymal-B-Play-v0 - Manager Based - **rsl_rl** (PPO), **skrl** (PPO) + - ``newton``, ``physx`` * - Isaac-Velocity-Flat-Anymal-C-Direct-v0 - - Direct - **rl_games** (PPO), **rsl_rl** (PPO), **skrl** (PPO) + - * - Isaac-Velocity-Flat-Anymal-C-v0 - Isaac-Velocity-Flat-Anymal-C-Play-v0 - Manager Based - **rsl_rl** (PPO), **rl_games** (PPO), **skrl** (PPO) + - ``newton``, ``physx`` * - Isaac-Velocity-Flat-Anymal-D-v0 - Isaac-Velocity-Flat-Anymal-D-Play-v0 - Manager Based - **rsl_rl** (PPO), **skrl** (PPO) + - ``newton``, ``physx`` * - Isaac-Velocity-Flat-Cassie-v0 - Isaac-Velocity-Flat-Cassie-Play-v0 - Manager Based - **rsl_rl** (PPO), **skrl** (PPO) + - * - Isaac-Velocity-Flat-Digit-v0 - Isaac-Velocity-Flat-Digit-Play-v0 - Manager Based - **rsl_rl** (PPO) + - ``newton``, ``physx`` * - Isaac-Velocity-Flat-G1-v0 - Isaac-Velocity-Flat-G1-Play-v0 - Manager Based - **rsl_rl** (PPO), **skrl** (PPO) + - ``newton``, ``physx`` * - Isaac-Velocity-Flat-H1-v0 - Isaac-Velocity-Flat-H1-Play-v0 - Manager Based - **rsl_rl** (PPO), **skrl** (PPO) + - ``newton``, ``physx`` * - Isaac-Velocity-Flat-Spot-v0 - Isaac-Velocity-Flat-Spot-Play-v0 - Manager Based - **rsl_rl** (PPO), **skrl** (PPO) + - ``newton``, ``physx`` * - Isaac-Velocity-Flat-Unitree-A1-v0 - Isaac-Velocity-Flat-Unitree-A1-Play-v0 - Manager Based - **rsl_rl** (PPO), **skrl** (PPO), **sb3** (PPO) + - ``newton``, ``physx`` * - Isaac-Velocity-Flat-Unitree-Go1-v0 - Isaac-Velocity-Flat-Unitree-Go1-Play-v0 - Manager Based - **rsl_rl** (PPO), **skrl** (PPO) + - ``newton``, ``physx`` * - Isaac-Velocity-Flat-Unitree-Go2-v0 - Isaac-Velocity-Flat-Unitree-Go2-Play-v0 - Manager Based - **rsl_rl** (PPO), **skrl** (PPO) + - ``newton``, ``physx`` * - Isaac-Velocity-Rough-Anymal-B-v0 - Isaac-Velocity-Rough-Anymal-B-Play-v0 - Manager Based - **rsl_rl** (PPO), **skrl** (PPO) + - ``newton``, ``physx`` * - Isaac-Velocity-Rough-Anymal-C-Direct-v0 - - Direct - **rl_games** (PPO), **rsl_rl** (PPO), **skrl** (PPO) + - * - Isaac-Velocity-Rough-Anymal-C-v0 - Isaac-Velocity-Rough-Anymal-C-Play-v0 - Manager Based - **rl_games** (PPO), **rsl_rl** (PPO), **skrl** (PPO) + - ``newton``, ``physx`` * - Isaac-Velocity-Rough-Anymal-D-v0 - Isaac-Velocity-Rough-Anymal-D-Play-v0 - Manager Based - **rsl_rl** (PPO), **skrl** (PPO) + - ``newton``, ``physx`` * - Isaac-Velocity-Rough-Cassie-v0 - Isaac-Velocity-Rough-Cassie-Play-v0 - Manager Based - **rsl_rl** (PPO), **skrl** (PPO) + - * - Isaac-Velocity-Rough-Digit-v0 - Isaac-Velocity-Rough-Digit-Play-v0 - Manager Based - **rsl_rl** (PPO) + - ``newton``, ``physx`` * - Isaac-Velocity-Rough-G1-v0 - Isaac-Velocity-Rough-G1-Play-v0 - Manager Based - **rsl_rl** (PPO), **skrl** (PPO) + - ``newton``, ``physx`` * - Isaac-Velocity-Rough-H1-v0 - Isaac-Velocity-Rough-H1-Play-v0 - Manager Based - **rsl_rl** (PPO), **skrl** (PPO) + - ``newton``, ``physx`` * - Isaac-Velocity-Rough-Unitree-A1-v0 - Isaac-Velocity-Rough-Unitree-A1-Play-v0 - Manager Based - **rsl_rl** (PPO), **skrl** (PPO), **sb3** (PPO) + - ``newton``, ``physx`` * - Isaac-Velocity-Rough-Unitree-Go1-v0 - Isaac-Velocity-Rough-Unitree-Go1-Play-v0 - Manager Based - **rsl_rl** (PPO), **skrl** (PPO) + - ``newton``, ``physx`` * - Isaac-Velocity-Rough-Unitree-Go2-v0 - Isaac-Velocity-Rough-Unitree-Go2-Play-v0 - Manager Based - **rsl_rl** (PPO), **skrl** (PPO) + - ``newton``, ``physx`` * - Isaac-Reach-OpenArm-Bi-v0 - Isaac-Reach-OpenArm-Bi-Play-v0 - Manager Based - **rsl_rl** (PPO), **rl_games** (PPO) + - * - Isaac-Reach-OpenArm-v0 - Isaac-Reach-OpenArm-Play-v0 - Manager Based - **rsl_rl** (PPO), **skrl** (PPO), **rl_games** (PPO) + - * - Isaac-Lift-Cube-OpenArm-v0 - Isaac-Lift-Cube-OpenArm-Play-v0 - Manager Based - **rsl_rl** (PPO), **rl_games** (PPO) + - * - Isaac-Open-Drawer-OpenArm-v0 - Isaac-Open-Drawer-OpenArm-Play-v0 - Manager Based - **rsl_rl** (PPO), **rl_games** (PPO) + - diff --git a/docs/source/setup/installation/cloud_installation.rst b/docs/source/setup/installation/cloud_installation.rst index cadecfce1c43..b6d9137680a3 100644 --- a/docs/source/setup/installation/cloud_installation.rst +++ b/docs/source/setup/installation/cloud_installation.rst @@ -2,15 +2,15 @@ Cloud Deployment ================ Isaac Lab can be run in various cloud infrastructures with the use of -`Isaac Automator `__. +`Isaac Automator `__ (v4). -Isaac Automator allows for quick deployment of Isaac Sim and Isaac Lab onto -the public clouds (AWS, GCP, Azure, and Alibaba Cloud are currently supported). -The result is a fully configured remote desktop cloud workstation, which can -be used for development and testing of Isaac Lab within minutes and on a budget. -Isaac Automator supports variety of GPU instances and stop-start functionality -to save on cloud costs and a variety of tools to aid the workflow -(such as uploading and downloading data, autorun, deployment management, etc). +Isaac Automator allows quick deployment of Isaac Sim, Isaac Lab, and Isaac Lab Arena +onto public clouds (AWS, GCP, Azure, and Alibaba Cloud are currently supported). +The result is a fully configured remote desktop cloud workstation (Isaac Workstation), +which can be used for development and testing of Isaac Lab within minutes and on a budget. +Isaac Automator supports a variety of GPU instances and stop/start functionality +to save on cloud costs, and provides tools to aid the workflow +(uploading and downloading data, autorun scripts, deployment management, etc.). System Requirements @@ -19,17 +19,16 @@ System Requirements Isaac Automator requires having ``docker`` pre-installed on the system. * To install Docker, please follow the instructions for your operating system on the - `Docker website`_. A minimum version of 26.0.0 for Docker Engine and 2.25.0 for Docker - compose are required to work with Isaac Automator. + `Docker website`_. * Follow the post-installation steps for Docker on the `post-installation steps`_ page. These steps allow you to run Docker without using ``sudo``. Installing Isaac Automator --------------------------- +--------------------------- -For the most update-to-date and complete installation instructions, please refer to -`Isaac Automator `__. +For the most up-to-date and complete installation instructions, please refer to +the `Isaac Automator README `__. To use Isaac Automator, first clone the repo: @@ -48,37 +47,15 @@ To use Isaac Automator, first clone the repo: git clone git@github.com:isaac-sim/IsaacAutomator.git -Isaac Automator requires obtaining a NGC API key. - -* Get access to the `Isaac Sim container`_ by joining the NVIDIA Developer Program credentials. -* Generate your `NGC API key`_ to access locked container images from NVIDIA GPU Cloud (NGC). - - * This step requires you to create an NGC account if you do not already have one. - * Once you have your generated API key, you need to log in to NGC - from the terminal. - - .. code:: bash - - docker login nvcr.io - - * For the username, enter ``$oauthtoken`` exactly as shown. It is a special username that is used to - authenticate with NGC. - - .. code:: text - - Username: $oauthtoken - Password: - - -Building the container +Building the Container ---------------------- -To run Isaac Automator, first build the Isaac Automator container: +Build the Isaac Automator container: .. tab-set:: :sync-group: os - .. tab-item:: :icon:`fa-brands fa-linux` Linux + .. tab-item:: :icon:`fa-brands fa-linux` Linux / macOS :sync: linux .. code-block:: bash @@ -90,149 +67,138 @@ To run Isaac Automator, first build the Isaac Automator container: .. code-block:: batch - docker build --platform linux/x86_64 -t isa . + docker build --platform linux/x86_64 -t isaac_automator . +This will build the Isaac Automator container and tag it as ``isaac_automator``. -This will build the Isaac Automator container and tag it as ``isa``. - -Running the Automator Commands +Deploying an Isaac Workstation ------------------------------ -First, enter the Automator container: - .. tab-set:: :sync-group: os - .. tab-item:: :icon:`fa-brands fa-linux` Linux + .. tab-item:: :icon:`fa-brands fa-linux` Linux / macOS :sync: linux + Enter the Automator container and run the deployment command: + .. code-block:: bash ./run + # inside container: + ./deploy-aws + + Alternatively, run it in one step: + + .. code-block:: bash + + ./run ./deploy-aws .. tab-item:: :icon:`fa-brands fa-windows` Windows :sync: windows .. code-block:: batch - docker run --platform linux/x86_64 -it --rm -v .:/app isa bash + docker run --platform linux/x86_64 -it --rm -v .:/app isaac_automator bash + :: inside container: + ./deploy-aws -Next, run the deployment script for your preferred cloud: +Replace ``deploy-aws`` with ``deploy-gcp``, ``deploy-azure``, or ``deploy-alicloud`` +for other cloud providers. .. note:: - The ``--isaaclab`` flag is used to specify the version of Isaac Lab to deploy. - The ``v3.0.0`` tag is the latest release of Isaac Lab. - -.. tab-set:: - :sync-group: cloud - - .. tab-item:: AWS - :sync: aws - - .. code-block:: bash - - ./deploy-aws --isaaclab v3.0.0 - - .. tab-item:: Azure - :sync: azure + The ``--isaaclab`` and ``--isaacsim`` flags accept any valid Git reference + to specify the version to deploy. Use ``--isaaclab no`` or ``--isaacsim no`` + to skip installation of the respective component. - .. code-block:: bash + .. code-block:: bash - ./deploy-azure --isaaclab v3.0.0 + ./deploy-aws --isaaclab v3.0.0 --isaacsim main - .. tab-item:: GCP - :sync: gcp +On the first run (or when credentials expire), you will be prompted to enter +your cloud credentials. Credentials are stored in ``state/`` and persist +across container restarts. Run ``./deploy- --help`` to see all available +options. - .. code-block:: bash +Key deployment options: - ./deploy-gcp --isaaclab v3.0.0 - - .. tab-item:: Alibaba Cloud - :sync: alicloud - - .. code-block:: bash +- ``--instance-type`` -- Cloud VM instance type. +- ``--isaacsim`` / ``--isaaclab`` / ``--isaaclab-arena`` -- Git ref for the version + to install, or ``no`` to skip. +- ``--existing`` -- What to do if a deployment already exists: ``ask`` (default), + ``repair``, ``modify``, ``replace``, or ``run_ansible``. +- ``--from-image`` -- Deploy from a pre-built VM image for faster provisioning + (AWS only at this time). - ./deploy-alicloud --isaaclab v3.0.0 +Connecting to the Isaac Workstation +----------------------------------- -Follow the prompts for entering information regarding the environment setup and credentials. -Once successful, instructions for connecting to the cloud instance will be available -in the terminal. The deployed Isaac Sim instances can be accessed via: +Deployed Isaac Workstations can be accessed via: -- SSH -- noVCN (browser-based VNC client) -- NoMachine (remote desktop client) +- **SSH**: ``./ssh `` +- **noVNC** (browser-based): ``./novnc `` +- **NoMachine** (remote desktop client) -Look for the connection instructions at the end of the deployment command output. -Additionally, this info is saved in ``state//info.txt`` file. - -For details on the credentials and setup required for each cloud, please visit the -`Isaac Automator `__ -page for more instructions. +Connection instructions are displayed at the end of the deployment command +output and saved in ``state//info.txt``. Running Isaac Lab on the Cloud ------------------------------ -Once connected to the cloud instance, the desktop will have an icon showing ``isaaclab.sh``. -Launch the ``isaaclab.sh`` executable, which will open a new Terminal. Within the terminal, -Isaac Lab commands can be executed in the same way as running locally. - -For example: - -.. tab-set:: - :sync-group: os - - .. tab-item:: :icon:`fa-brands fa-linux` Linux - :sync: linux +Isaac Lab is installed from source on the deployed workstation at ``~/IsaacLab``. +To run Isaac Lab commands, open a terminal on the workstation: - .. code-block:: bash +.. code-block:: bash - ./isaaclab.sh -p scripts/reinforcement_learning/rl_games/train.py --task=Isaac-Cartpole-v0 + ~/IsaacLab/isaaclab.sh -p scripts/reinforcement_learning/rsl_rl/train.py \ + --task=Isaac-Cartpole-Direct-v0 --headless - .. tab-item:: :icon:`fa-brands fa-windows` Windows - :sync: windows - .. code-block:: batch +Pausing and Resuming +-------------------- - isaaclab.bat -p scripts/reinforcement_learning/rl_games/train.py --task=Isaac-Cartpole-v0 +You can stop and restart instances to save on cloud costs: +.. code-block:: bash -Destroying a Deployment ------------------------ + # inside the Automator container: + ./stop + ./start -To save costs, deployments can be destroyed when not being used. -This can be done from within the Automator container. +Use ``./start --quick`` to skip full Ansible provisioning +and only run the autorun script. -Enter the Automator container with the command described in the previous section: -.. tab-set:: - :sync-group: os +Uploading and Downloading Data +------------------------------ - .. tab-item:: :icon:`fa-brands fa-linux` Linux - :sync: linux +.. code-block:: bash - .. code-block:: bash + # upload local uploads/ folder to the instance + ./upload - ./run + # download results from the instance to local results/ folder + ./download - .. tab-item:: :icon:`fa-brands fa-windows` Windows - :sync: windows - .. code-block:: batch +Destroying a Deployment +----------------------- - docker run --platform linux/x86_64 -it --rm -v .:/app isa bash +To save costs, destroy deployments when no longer needed: +.. code-block:: bash -To destroy a deployment, run the following command from within the container: + # inside the Automator container: + ./destroy -.. code:: bash +.. note:: - ./destroy + Deployment metadata is stored in the ``state/`` directory. Do not delete this + directory, as it is required for managing deployments. -.. _`Docker website`: https://docs.docker.com/desktop/install/linux-install/ +.. _`Docker website`: https://docs.docker.com/engine/install/ .. _`post-installation steps`: https://docs.docker.com/engine/install/linux-postinstall/ -.. _`Isaac Sim container`: https://catalog.ngc.nvidia.com/orgs/nvidia/containers/isaac-sim -.. _`NGC API key`: https://docs.nvidia.com/ngc/gpu-cloud/ngc-user-guide/index.html#generating-api-key diff --git a/docs/source/setup/installation/include/src_python_virtual_env.rst b/docs/source/setup/installation/include/src_python_virtual_env.rst index 617e29ace75a..4ca31fafb17a 100644 --- a/docs/source/setup/installation/include/src_python_virtual_env.rst +++ b/docs/source/setup/installation/include/src_python_virtual_env.rst @@ -62,7 +62,7 @@ instead of *./isaaclab.sh -p* or *isaaclab.bat -p*. .. warning:: Windows support for UV is currently unavailable. Please check - `issue #3483 `_ to track progress. + `issue #3438 `_ to track progress. .. tab-item:: Conda Environment @@ -103,7 +103,7 @@ instead of *./isaaclab.sh -p* or *isaaclab.bat -p*. .. code:: batch :: Activate environment - conda activate env_isaaclab # or "conda activate my_env" + conda activate env_isaaclab :: or "conda activate my_env" Once you are in the virtual environment, you do not need to use ``./isaaclab.sh -p`` or ``isaaclab.bat -p`` to run python scripts. You can use the default python executable in your diff --git a/docs/source/setup/installation/isaaclab_pip_installation.rst b/docs/source/setup/installation/isaaclab_pip_installation.rst index d3070fd51f50..1d98f1536971 100644 --- a/docs/source/setup/installation/isaaclab_pip_installation.rst +++ b/docs/source/setup/installation/isaaclab_pip_installation.rst @@ -61,7 +61,7 @@ Isaac Lab sub-packages: uv pip install isaaclab==3.0.0 # specific version # Isaac Lab + Isaac Sim - uv pip install "isaaclab[isaacsim]" --index-strategy unsafe-best-match --prerelease=allow + uv pip install "isaaclab[isaacsim]" --extra-index-url https://pypi.nvidia.com --index-strategy unsafe-best-match --prerelease=allow # Isaac Lab + specific sub-package(s) # Note: flags above are only needed when installing the isaacsim extra @@ -69,7 +69,7 @@ Isaac Lab sub-packages: uv pip install "isaaclab[rl,tasks]" # Isaac Lab + Isaac Sim + all sub-packages - uv pip install "isaaclab[isaacsim,all]" --index-strategy unsafe-best-match --prerelease=allow + uv pip install "isaaclab[isaacsim,all]" --extra-index-url https://pypi.nvidia.com --index-strategy unsafe-best-match --prerelease=allow .. tab-item:: pip @@ -170,8 +170,9 @@ Installing dependencies When using a conda environment, the preload is set up via the conda activation hook. -- If you want to use ``rl_games`` for training and inferencing, install - its Python 3.11+ enabled fork: +- If you want to use ``rl_games`` for training and inferencing **and did not + install the** ``rl`` **extra above**, install its Python 3.11+ enabled fork + manually: .. code-block:: none diff --git a/docs/source/setup/installation/source_installation.rst b/docs/source/setup/installation/source_installation.rst index c697c1dd2054..ce575e48bc7b 100644 --- a/docs/source/setup/installation/source_installation.rst +++ b/docs/source/setup/installation/source_installation.rst @@ -78,7 +78,7 @@ variables to your terminal for the remaining of the installation instructions: .. code:: bash # Isaac Sim root directory - export ISAACSIM_PATH="${pwd}/_build/linux-x86_64/release" + export ISAACSIM_PATH="${PWD}/_build/linux-x86_64/release" # Isaac Sim python executable export ISAACSIM_PYTHON_EXE="${ISAACSIM_PATH}/python.sh" diff --git a/docs/source/setup/quick_installation.rst b/docs/source/setup/quick_installation.rst index 3a4b60cb3317..36ed9208fd39 100644 --- a/docs/source/setup/quick_installation.rst +++ b/docs/source/setup/quick_installation.rst @@ -15,8 +15,8 @@ Quick Installation cd IsaacLab # Create environment and install - uv venv .venv --python 3.12 - source .venv/bin/activate + uv venv --python 3.12 --seed env_isaaclab + source env_isaaclab/bin/activate ./isaaclab.sh -i # Run training (Newton backend, 16 envs) diff --git a/docs/source/setup/quickstart.rst b/docs/source/setup/quickstart.rst index 4ebdcd024046..39e7ec4cb932 100644 --- a/docs/source/setup/quickstart.rst +++ b/docs/source/setup/quickstart.rst @@ -68,11 +68,31 @@ package manager. To begin, create a virtual environment: conda activate env_isaaclab -Next, install a CUDA-enabled PyTorch build. +Next, install a CUDA-enabled PyTorch build that matches your system architecture. - .. code-block:: bash +.. tab-set:: + :sync-group: pip-platform + + .. tab-item:: :icon:`fa-brands fa-linux` Linux (x86_64) + :sync: linux-x86_64 + + .. code-block:: bash + + uv pip install -U torch==2.10.0 torchvision==0.25.0 --index-url https://download.pytorch.org/whl/cu128 + + .. tab-item:: :icon:`fa-brands fa-windows` Windows (x86_64) + :sync: windows-x86_64 + + .. code-block:: bash + + uv pip install -U torch==2.10.0 torchvision==0.25.0 --index-url https://download.pytorch.org/whl/cu128 + + .. tab-item:: :icon:`fa-brands fa-linux` Linux (aarch64) + :sync: linux-aarch64 + + .. code-block:: bash - uv pip install -U torch==2.10.0 torchvision==0.25.0 --index-url https://download.pytorch.org/whl/cu128 + uv pip install -U torch==2.10.0 torchvision==0.25.0 --index-url https://download.pytorch.org/whl/cu130 Before we can install Isaac Sim, we need to make sure pip is updated. To update pip, run diff --git a/docs/source/testing/index.rst b/docs/source/testing/index.rst index ae8494a3ec89..4d875d982797 100644 --- a/docs/source/testing/index.rst +++ b/docs/source/testing/index.rst @@ -1,7 +1,7 @@ .. _testing: Testing -======= +======== This section covers testing utilities and patterns for Isaac Lab development. 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/isaaclab.bat b/isaaclab.bat index 077e9a4b1abd..1d8fb8275467 100644 --- a/isaaclab.bat +++ b/isaaclab.bat @@ -26,6 +26,13 @@ if defined VIRTUAL_ENV ( rem Add source/isaaclab to PYTHONPATH so we can import isaaclab.cli. set "PYTHONPATH=%ISAACLAB_PATH%\source\isaaclab;%PYTHONPATH%" +rem If a local Isaac Sim binary is present, source its env setup so that +rem PYTHONPATH/PATH/EXP_PATH are correct without depending on a conda +rem activate.d hook (those don't fire under e.g. `conda run` on Windows). +if exist "%ISAACLAB_PATH%\_isaac_sim\setup_conda_env.bat" ( + call "%ISAACLAB_PATH%\_isaac_sim\setup_conda_env.bat" >NUL +) + rem Execute CLI. "%python_exe%" -c "from isaaclab.cli import cli; cli()" %* diff --git a/isaaclab.sh b/isaaclab.sh index 8d535a10b307..d4042353e88f 100755 --- a/isaaclab.sh +++ b/isaaclab.sh @@ -28,5 +28,13 @@ fi # Add source/isaaclab to PYTHONPATH so we can import isaaclab.cli. export PYTHONPATH="$ISAACLAB_PATH/source/isaaclab:$PYTHONPATH" +# If a local Isaac Sim binary is present, source its env setup so that +# PYTHONPATH/PATH/EXP_PATH are correct without depending on a conda +# activate.d hook (those don't fire reliably under e.g. `conda run`). +if [ -f "$ISAACLAB_PATH/_isaac_sim/setup_conda_env.sh" ]; then + # shellcheck disable=SC1091 + . "$ISAACLAB_PATH/_isaac_sim/setup_conda_env.sh" >/dev/null 2>&1 || true +fi + # Execute CLI. exec "$python_exe" -c "from isaaclab.cli import cli; cli()" "$@" 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 43c6eb6911e0..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"], @@ -127,13 +127,13 @@ def run_simulator(sim: sim_utils.SimulationContext, scene: InteractiveScene): # print information from the sensors print("-------------------------------") print(scene["ray_caster"]) - print("Ray cast hit results: ", scene["ray_caster"].data.ray_hits_w) + print("Ray cast hit results: ", wp.to_torch(scene["ray_caster"].data.ray_hits_w)) if not triggered: if countdown > 0: countdown -= 1 continue - data = scene["ray_caster"].data.ray_hits_w.cpu().numpy() + data = wp.to_torch(scene["ray_caster"].data.ray_hits_w).cpu().numpy() np.save("cast_data.npy", data) triggered = True else: diff --git a/scripts/environments/teleoperation/teleop_se3_agent.py b/scripts/environments/teleoperation/teleop_se3_agent.py index 897eb159e86e..cdd5c104c44f 100644 --- a/scripts/environments/teleoperation/teleop_se3_agent.py +++ b/scripts/environments/teleoperation/teleop_se3_agent.py @@ -218,7 +218,7 @@ def stop_teleoperation() -> None: try: if use_isaac_teleop: - from isaaclab_teleop import create_isaac_teleop_device + from isaaclab_teleop import create_isaac_teleop_device, poll_control_events teleop_interface = create_isaac_teleop_device( env_cfg.isaac_teleop, @@ -297,6 +297,13 @@ def run_loop(): # get device command action = teleop_interface.advance() + if use_isaac_teleop: + ctrl = poll_control_events(teleop_interface) + if ctrl.is_active is not None: + teleoperation_active = ctrl.is_active + if ctrl.should_reset: + should_reset_recording_instance = True + # action is None when IsaacTeleop session hasn't started yet # (e.g. waiting for user to click "Start AR") if action is None: diff --git a/scripts/tools/record_demos.py b/scripts/tools/record_demos.py index bd318b7a2625..75df9e0ee92a 100644 --- a/scripts/tools/record_demos.py +++ b/scripts/tools/record_demos.py @@ -406,26 +406,34 @@ def process_success_condition(env: gym.Env, success_term: object | None, success def handle_reset( - env: gym.Env, success_step_count: int, instruction_display: InstructionDisplay, label_text: str + env: gym.Env, + success_step_count: int, + instruction_display: InstructionDisplay, + label_text: str, + teleop_interface: object | None = None, ) -> int: """Handle resetting the environment. - Resets the environment, recorder manager, and related state variables. - Updates the instruction display with current status. + Resets the environment, recorder manager, teleop device, and related + state variables. Updates the instruction display with current status. Args: - env: The environment instance to reset - success_step_count: Current count of consecutive successful steps - instruction_display: The display object to update - label_text: Text to display showing current recording status + env: The environment instance to reset. + success_step_count: Current count of consecutive successful steps. + instruction_display: The display object to update. + label_text: Text to display showing current recording status. + teleop_interface: Optional teleop device to reset (resets XR anchor + and retargeter cross-step state). Returns: - int: Reset success step count (0) + Reset success step count (0). """ print("Resetting environment...") env.sim.reset() env.recorder_manager.reset() env.reset() + if teleop_interface is not None and hasattr(teleop_interface, "reset"): + teleop_interface.reset() success_step_count = 0 instruction_display.show_demo(label_text) return success_step_count @@ -476,7 +484,9 @@ def stop_recording_instance(): running_recording_instance = False print("Recording paused") - # Set up teleoperation callbacks + # Set up teleoperation callbacks. For IsaacTeleop the primary control + # path is poll_control_events(); these callbacks are bridged automatically + # and also serve native (keyboard / spacemouse) devices. teleoperation_callbacks = { "R": reset_recording_instance, "START": start_recording_instance, @@ -485,7 +495,6 @@ def stop_recording_instance(): } teleop_interface = setup_teleop_device(teleoperation_callbacks, use_isaac_teleop) - teleop_interface.add_callback("R", reset_recording_instance) label_text = f"Recorded {current_recorded_demo_count} successful demonstrations." instruction_display = setup_ui(label_text, env) @@ -504,10 +513,21 @@ def inner_loop(): stack_name = "IsaacTeleop" if use_isaac_teleop else "native" print(f"{stack_name} recording started.") + if use_isaac_teleop: + from isaaclab_teleop import poll_control_events + with contextlib.suppress(KeyboardInterrupt), torch.inference_mode(): while simulation_app.is_running(): # Get teleop command (may be None while waiting for session start) action = teleop_interface.advance() + + if use_isaac_teleop: + ctrl = poll_control_events(teleop_interface) + if ctrl.is_active is not None: + running_recording_instance = ctrl.is_active + if ctrl.should_reset: + should_reset_recording_instance = True + if action is None: env.sim.render() continue @@ -558,7 +578,9 @@ def inner_loop(): # Handle reset if requested if should_reset_recording_instance: - success_step_count = handle_reset(env, success_step_count, instruction_display, label_text) + success_step_count = handle_reset( + env, success_step_count, instruction_display, label_text, teleop_interface + ) should_reset_recording_instance = False # Check if simulation is stopped diff --git a/scripts/tutorials/04_sensors/add_sensors_on_robot.py b/scripts/tutorials/04_sensors/add_sensors_on_robot.py index 31f9a2bcefcb..f5e3a19c0bec 100644 --- a/scripts/tutorials/04_sensors/add_sensors_on_robot.py +++ b/scripts/tutorials/04_sensors/add_sensors_on_robot.py @@ -150,7 +150,10 @@ def run_simulator(sim: sim_utils.SimulationContext, scene: InteractiveScene): print("Received shape of depth image: ", scene["camera"].data.output["distance_to_image_plane"].shape) print("-------------------------------") print(scene["height_scanner"]) - print("Received max height value: ", torch.max(scene["height_scanner"].data.ray_hits_w[..., -1]).item()) + print( + "Received max height value: ", + torch.max(wp.to_torch(scene["height_scanner"].data.ray_hits_w)[..., -1]).item(), + ) print("-------------------------------") print(scene["contact_forces"]) print("Received max contact force of: ", torch.max(scene["contact_forces"].data.net_forces_w).item()) diff --git a/scripts/tutorials/04_sensors/run_ray_caster.py b/scripts/tutorials/04_sensors/run_ray_caster.py index 3e46ef1a08fd..ff66ff9a0fd2 100644 --- a/scripts/tutorials/04_sensors/run_ray_caster.py +++ b/scripts/tutorials/04_sensors/run_ray_caster.py @@ -120,7 +120,7 @@ def run_simulator(sim: sim_utils.SimulationContext, scene_entities: dict): # Update the ray-caster with Timer( f"Ray-caster update with {4} x {ray_caster.num_rays} rays with max height of" - f" {torch.max(ray_caster.data.pos_w).item():.2f}" + f" {torch.max(wp.to_torch(ray_caster.data.pos_w)).item():.2f}" ): ray_caster.update(dt=sim.get_physics_dt(), force_recompute=True) # Update counter diff --git a/source/isaaclab/config/extension.toml b/source/isaaclab/config/extension.toml index 85b9e265b1b7..3086a2b93c88 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.7" +version = "4.6.12" # Description title = "Isaac Lab framework for Robot Learning" diff --git a/source/isaaclab/docs/CHANGELOG.rst b/source/isaaclab/docs/CHANGELOG.rst index a6f25d38ee40..d5af850f2412 100644 --- a/source/isaaclab/docs/CHANGELOG.rst +++ b/source/isaaclab/docs/CHANGELOG.rst @@ -1,6 +1,153 @@ Changelog --------- +4.6.12 (2026-04-23) +~~~~~~~~~~~~~~~~~~~ + +Added +^^^^^ + +* Added caching to :func:`~isaaclab.utils.string.resolve_matching_names`, + avoiding repeated regex matching across ``find_bodies``, ``find_joints``, + and related calls. + + +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) +~~~~~~~~~~~~~~~~~~~ + +Added +^^^^^ + +* Added :meth:`~isaaclab.utils.wrench_composer.WrenchComposer.add_raw_buffers_from` to merge one composer's raw + input buffers into another. + +Changed +^^^^^^^ + +* Refactored :class:`~isaaclab.utils.wrench_composer.WrenchComposer` to a dual-buffer architecture with separate + global (world-frame) and local (body-frame) buffers. A new + :meth:`~isaaclab.utils.wrench_composer.WrenchComposer.compose_to_body_frame` method rotates global forces/torques + into the body frame at apply time using the current body orientation, then sums with local forces/torques. + +Deprecated +^^^^^^^^^^ + +* Deprecated :attr:`~isaaclab.utils.wrench_composer.WrenchComposer.composed_force` and + :attr:`~isaaclab.utils.wrench_composer.WrenchComposer.composed_torque` in favor of + :attr:`~isaaclab.utils.wrench_composer.WrenchComposer.out_force_b` and + :attr:`~isaaclab.utils.wrench_composer.WrenchComposer.out_torque_b`. + +Fixed +^^^^^ + +* Fixed :class:`~isaaclab.utils.wrench_composer.WrenchComposer` not correctly updating the composed torque from global + positional forces when the body moves. +* Fixed :meth:`~isaaclab.utils.wrench_composer.WrenchComposer.reset` not clearing the ``_active`` flag when called + with ``slice(None)``. +* Fixed :class:`~isaaclab.utils.wrench_composer.WrenchComposer` producing spurious torque when global forces are + applied without explicit positions. +* Fixed ``set_external_force_and_torque`` wiping forces from non-resetting environments during partial + episode resets by using ``reset(env_ids)`` + ``add_forces_and_torques`` instead of ``set_forces_and_torques``. + + +4.6.9 (2026-04-22) +~~~~~~~~~~~~~~~~~~ + +Changed +^^^^^^^ + +* Converted all four ray caster sensor classes (:class:`~isaaclab.sensors.RayCaster`, + :class:`~isaaclab.sensors.RayCasterCamera`, :class:`~isaaclab.sensors.MultiMeshRayCaster`, + :class:`~isaaclab.sensors.MultiMeshRayCasterCamera`) to launch Warp kernels directly via + ``wp.launch`` instead of going through Python-level torch wrappers. A new + :mod:`~isaaclab.sensors.ray_caster.kernels` module contains all sensor-specific kernels. + All intermediate ray buffers are now Warp-owned with zero-copy torch views, eliminating + per-step allocations. The existing :func:`~isaaclab.utils.warp.kernels.raycast_dynamic_meshes_kernel` + gained an ``env_mask`` parameter to support partial environment updates natively. A new + :func:`~isaaclab.utils.warp.kernels.raycast_mesh_masked_kernel` was added to + :mod:`~isaaclab.utils.warp.kernels` as the general-purpose masked single-mesh variant, + with ``return_distance`` and ``return_normal`` flags matching the design of + :func:`~isaaclab.utils.warp.kernels.raycast_mesh_kernel`. + + **Breaking change** — :attr:`~isaaclab.sensors.RayCasterData.pos_w`, + :attr:`~isaaclab.sensors.RayCasterData.quat_w`, and + :attr:`~isaaclab.sensors.RayCasterData.ray_hits_w` now return :class:`wp.array` + instead of :class:`torch.Tensor`. Call-sites that previously accessed these as tensors + must wrap the result with :func:`wp.to_torch`: + + .. code-block:: python + + # Before + hits = sensor.data.ray_hits_w # torch.Tensor (old) + # After + hits = wp.to_torch(sensor.data.ray_hits_w) # torch.Tensor (zero-copy view) + +* Changed the :attr:`~isaaclab.sensors.RayCaster.meshes` class variable cache key from + ``prim_path`` to a ``(prim_path, device)`` tuple so that meshes built on one device + (e.g. CPU) are not reused by a sensor running on another device (e.g. CUDA). + + **Breaking change** — callers that read or write :attr:`~isaaclab.sensors.RayCaster.meshes` + directly must update the key: + + .. code-block:: python + + # Before + wp_mesh = RayCaster.meshes[prim_path] + # After + wp_mesh = RayCaster.meshes[(prim_path, device)] + +Fixed +^^^^^ + +* Fixed frame composition in :meth:`~isaaclab.sensors.MultiMeshRayCaster._update_mesh_transforms` + which used simple subtraction instead of proper frame decomposition when applying mesh offsets. + With non-identity orientation offsets, tracked mesh positions were incorrect, causing raycasts to + miss or hit wrong surfaces. The method now uses :func:`~isaaclab.utils.math.combine_frame_transforms`. + + +4.6.8 (2026-04-21) +~~~~~~~~~~~~~~~~~~ + +Changed +^^^^^^^ + +* Pinned ``mujoco`` and ``mujoco-warp`` to ``3.6.0`` to align with the Newton library. + + 4.6.7 (2026-04-20) ~~~~~~~~~~~~~~~~~~ @@ -279,6 +426,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/app/app_launcher.py b/source/isaaclab/isaaclab/app/app_launcher.py index df6f7cb79e18..c9e18b54f54d 100644 --- a/source/isaaclab/isaaclab/app/app_launcher.py +++ b/source/isaaclab/isaaclab/app/app_launcher.py @@ -49,57 +49,6 @@ def __call__(self, parser, namespace, values, option_string=None): setattr(namespace, f"{self.dest}_explicit", True) -def _parse_visualizer_csv(value: str) -> list[str]: - """Parse visualizer list from a single comma-delimited CLI token.""" - valid = {"kit", "newton", "rerun", "viser", "none"} - token = (value or "").strip() - if not token: - raise argparse.ArgumentTypeError( - "Invalid --visualizer value: empty string. Use a comma-separated list, e.g. --viz kit,newton." - ) - if " " in token: - raise argparse.ArgumentTypeError( - "Invalid --visualizer value: spaces are not allowed. " - "Use a comma-separated list without spaces, e.g. --viz kit,newton,rerun,viser." - ) - - names = [item.strip().lower() for item in token.split(",")] - if any(not name for name in names): - raise argparse.ArgumentTypeError( - "Invalid --visualizer value: empty visualizer entry detected. " - "Use a comma-separated list without empty items." - ) - invalid = [name for name in names if name not in valid] - if invalid: - raise argparse.ArgumentTypeError( - f"Invalid --visualizer value(s): {', '.join(invalid)}. Valid options: {', '.join(sorted(valid))}." - ) - # De-duplicate while preserving order. - return list(dict.fromkeys(names)) - - -def _normalize_visualizer_intent(intent: Any) -> tuple[bool, bool]: - """Normalize and validate upstream config visualizer intent payload. - - The expected schema is: - ``{"has_any_visualizers": bool, "has_kit_visualizer": bool}``. - """ - if intent is None: - return False, False - if not isinstance(intent, dict): - raise ValueError("Invalid value for `visualizer_intent`: expected dict or None.") - - has_any = intent.get("has_any_visualizers", False) - has_kit = intent.get("has_kit_visualizer", False) - if not isinstance(has_any, bool) or not isinstance(has_kit, bool): - raise ValueError( - "Invalid `visualizer_intent` values: expected booleans for `has_any_visualizers` and `has_kit_visualizer`." - ) - if has_kit and not has_any: - raise ValueError("Invalid `visualizer_intent`: `has_kit_visualizer=True` requires `has_any_visualizers=True`.") - return has_any, has_kit - - class ExplicitTrueAction(argparse.Action): """Custom action to track explicit use of boolean flags.""" @@ -133,6 +82,96 @@ class AppLauncher: """ + @staticmethod + def sync_visualizer_cli_settings_to_carb(launcher_args: dict) -> None: + """Write visualizer CLI selection and ``--max_visible_envs`` to carb settings. + + Callers may set ``visualizer_explicit`` / ``visualizer_disable_all`` when those values + were resolved elsewhere (e.g. :class:`AppLauncher` strips flags from *launcher_args*). + Otherwise ``disable_all`` is inferred from ``"none"`` in ``visualizer``. + + Also used when Kit is skipped (see :mod:`isaaclab_tasks.utils.sim_launcher`). + """ + visualizers = launcher_args.get("visualizer") + + if "max_visible_envs" in launcher_args: + v = launcher_args["max_visible_envs"] + if v is not None and int(v) < 0: + raise ValueError(f"Invalid value for --max_visible_envs: {v}. Expected non-negative int.") + + cli_explicit = bool(launcher_args.get("visualizer_explicit", False)) + if "visualizer_disable_all" in launcher_args: + cli_disable_all = bool(launcher_args["visualizer_disable_all"]) + else: + cli_disable_all = bool(cli_explicit) and visualizers is not None and "none" in visualizers + + with contextlib.suppress(Exception): + visualizer_str = " ".join(visualizers) if visualizers else "" + settings = get_settings_manager() + settings.set_string("/isaaclab/visualizer/types", visualizer_str) + settings.set_bool("/isaaclab/visualizer/explicit", cli_explicit) + settings.set_bool("/isaaclab/visualizer/disable_all", cli_disable_all) + + # Sentinel: ``-1`` means ``--max_visible_envs`` was not passed (see ``SimulationContext``). + if "max_visible_envs" in launcher_args: + settings.set_int("/isaaclab/visualizer/max_visible_envs", int(launcher_args["max_visible_envs"])) + else: + settings.set_int("/isaaclab/visualizer/max_visible_envs", -1) + + @staticmethod + def _parse_visualizer_csv(value: str) -> list[str]: + """Parse visualizer list from a single comma-delimited CLI token.""" + valid = {"kit", "newton", "rerun", "viser", "none"} + token = (value or "").strip() + if not token: + raise argparse.ArgumentTypeError( + "Invalid --visualizer value: empty string. Use a comma-separated list, e.g. --viz kit,newton." + ) + if " " in token: + raise argparse.ArgumentTypeError( + "Invalid --visualizer value: spaces are not allowed. " + "Use a comma-separated list without spaces, e.g. --viz kit,newton,rerun,viser." + ) + + names = [item.strip().lower() for item in token.split(",")] + if any(not name for name in names): + raise argparse.ArgumentTypeError( + "Invalid --visualizer value: empty visualizer entry detected. " + "Use a comma-separated list without empty items." + ) + invalid = [name for name in names if name not in valid] + if invalid: + raise argparse.ArgumentTypeError( + f"Invalid --visualizer value(s): {', '.join(invalid)}. Valid options: {', '.join(sorted(valid))}." + ) + # De-duplicate while preserving order. + return list(dict.fromkeys(names)) + + @staticmethod + def _normalize_visualizer_intent(intent: Any) -> tuple[bool, bool]: + """Normalize and validate upstream config visualizer intent payload. + + The expected schema is: + ``{"has_any_visualizers": bool, "has_kit_visualizer": bool}``. + """ + if intent is None: + return False, False + if not isinstance(intent, dict): + raise ValueError("Invalid value for `visualizer_intent`: expected dict or None.") + + has_any = intent.get("has_any_visualizers", False) + has_kit = intent.get("has_kit_visualizer", False) + if not isinstance(has_any, bool) or not isinstance(has_kit, bool): + raise ValueError( + "Invalid `visualizer_intent` values: expected booleans for `has_any_visualizers` and " + "`has_kit_visualizer`." + ) + if has_kit and not has_any: + raise ValueError( + "Invalid `visualizer_intent`: `has_kit_visualizer=True` requires `has_any_visualizers=True`." + ) + return has_any, has_kit + def __init__(self, launcher_args: argparse.Namespace | dict | None = None, **kwargs): """Create a `SimulationApp`_ instance based on the input settings. @@ -188,7 +227,7 @@ def __init__(self, launcher_args: argparse.Namespace | dict | None = None, **kwa self._livestream: Literal[0, 1, 2] # 0: Disabled, 1: WebRTC public, 2: WebRTC private self._offscreen_render: bool # 0: Disabled, 1: Enabled self._sim_experience_file: str # Experience file to load - self._visualizer_max_worlds: int | None # Optional max worlds override for Newton-based visualizers + self._video_enabled: bool # Whether --video recording is enabled # Exposed to train scripts self.device_id: int # device ID for GPU simulation (defaults to 0) @@ -329,10 +368,9 @@ def add_app_launcher_args(parser: argparse.ArgumentParser) -> None: - Multiple visualizers can be specified as a comma-delimited list: ``--viz rerun,newton,viser``. - * ``visualizer_max_worlds`` (int | None): Optional global override for the maximum number of worlds - rendered in Newton-based visualizers (newton, rerun, viser). If omitted, each visualizer uses its - config default. - + * ``max_visible_envs`` (int | None): Optional global override for partial visualization by capping + how many environments are shown in the visualizers. + More partial visualization configuration fields are available in the ``VisualizerCfg`` class. .. _`WebRTC`: https://docs.isaacsim.omniverse.nvidia.com/latest/installation/manual_livestream_clients.html#isaac-sim-short-webrtc-streaming-client @@ -413,7 +451,7 @@ def add_app_launcher_args(parser: argparse.ArgumentParser) -> None: arg_group.add_argument( "--visualizer", "--viz", - type=_parse_visualizer_csv, + type=AppLauncher._parse_visualizer_csv, action=ExplicitAction, default=None, help="Visualizer backends to enable as CSV (e.g., kit,newton,rerun,viser).", @@ -484,13 +522,10 @@ def add_app_launcher_args(parser: argparse.ArgumentParser) -> None: ), ) arg_group.add_argument( - "--visualizer_max_worlds", + "--max_visible_envs", type=int, - default=AppLauncher._APPLAUNCHER_CFG_INFO["visualizer_max_worlds"][1], - help=( - "Optional global max worlds override for Newton-based visualizers (newton/rerun/viser). " - "If omitted, visualizer config defaults are used." - ), + default=argparse.SUPPRESS, + help=("When set, caps the nums of envs shown in the launched visualizers."), ) # special flag for backwards compatibility @@ -512,7 +547,7 @@ def add_app_launcher_args(parser: argparse.ArgumentParser) -> None: "device": ([str], "cuda:0"), "experience": ([str], ""), "rendering_mode": ([str], "balanced"), - "visualizer_max_worlds": ([int, type(None)], None), + "max_visible_envs": ([int, type(None)], None), } """A dictionary of arguments added manually by the :meth:`AppLauncher.add_app_launcher_args` method. @@ -782,7 +817,9 @@ def _resolve_headless_settings(self, launcher_args: dict, livestream_arg: int, l def _resolve_visualizer_settings(self, launcher_args: dict) -> None: """Resolve visualizer CLI semantics and normalize selection.""" raw_visualizers = launcher_args.get("visualizer") - cfg_has_any, cfg_has_kit = _normalize_visualizer_intent(launcher_args.pop("visualizer_intent", None)) + cfg_has_any, cfg_has_kit = AppLauncher._normalize_visualizer_intent( + launcher_args.pop("visualizer_intent", None) + ) self._cfg_has_any_visualizers = cfg_has_any self._cfg_has_kit_visualizer = cfg_has_kit visualizer_explicit = bool(launcher_args.pop("visualizer_explicit", False)) @@ -792,7 +829,7 @@ def _resolve_visualizer_settings(self, launcher_args: dict) -> None: visualizer_types: list[str] = [] if raw_visualizers is not None: if isinstance(raw_visualizers, str): - visualizer_types = _parse_visualizer_csv(raw_visualizers) + visualizer_types = AppLauncher._parse_visualizer_csv(raw_visualizers) else: visualizer_types = [str(v).strip().lower() for v in raw_visualizers if str(v).strip()] @@ -858,12 +895,13 @@ def _resolve_xr_settings(self, launcher_args: dict): def _resolve_viewport_settings(self, launcher_args: dict): """Resolve viewport related settings.""" + self._video_enabled = bool(launcher_args.get("video", False)) # Check if we can disable the viewport to improve performance # This should only happen if we are running headless and do not require livestreaming or video recording # This is different from offscreen_render because this only affects the default viewport and # not other render-products in the scene self._render_viewport = True - if self._headless and not self._livestream and not launcher_args.get("video", False): + if self._headless and not self._livestream and not self._video_enabled: self._render_viewport = False # hide_ui flag @@ -1085,6 +1123,8 @@ def _load_extensions(self): # (no Kit GUI) the AR profile must be enabled programmatically so that # the OpenXR session starts without user interaction settings.set_bool("/isaaclab/xr/auto_start", self._headless and self._xr) + # set setting to indicate video recording mode + settings.set_bool("/isaaclab/video/enabled", self._video_enabled) # set setting to indicate no RTX sensors are used (set to True when RTX sensor is created) settings.set_bool("/isaaclab/render/rtx_sensors", False) @@ -1148,28 +1188,14 @@ def _set_animation_recording_settings(self, launcher_args: dict) -> None: settings.set_float("/isaaclab/anim_recording/stop_time", stop_time) def _set_visualizer_settings(self, launcher_args: dict) -> None: - """Store visualizer selection and max-worlds override in settings.""" - visualizers = launcher_args.get("visualizer") - visualizer_max_worlds = launcher_args.get("visualizer_max_worlds") - - if visualizer_max_worlds is not None and visualizer_max_worlds < 0: - raise ValueError( - f"Invalid value for --visualizer_max_worlds: {visualizer_max_worlds}. Expected non-negative int." - ) - - with contextlib.suppress(Exception): - visualizer_str = " ".join(visualizers) if visualizers else "" - settings = get_settings_manager() - cli_visualizer_explicit = getattr(self, "_cli_visualizer_explicit", False) - cli_visualizer_disable_all = getattr(self, "_cli_visualizer_disable_all", False) - settings.set_string("/isaaclab/visualizer/types", visualizer_str) - settings.set_bool("/isaaclab/visualizer/explicit", cli_visualizer_explicit) - settings.set_bool("/isaaclab/visualizer/disable_all", cli_visualizer_disable_all) - # Store as int setting where -1 means "use per-visualizer defaults". - if visualizer_max_worlds is None: - settings.set_int("/isaaclab/visualizer/max_worlds", -1) - else: - settings.set_int("/isaaclab/visualizer/max_worlds", int(visualizer_max_worlds)) + """Persist visualizer CLI flags and ``max_visible_envs`` override for :class:`SimulationContext`.""" + AppLauncher.sync_visualizer_cli_settings_to_carb( + { + **launcher_args, + "visualizer_explicit": getattr(self, "_cli_visualizer_explicit", False), + "visualizer_disable_all": getattr(self, "_cli_visualizer_disable_all", False), + } + ) def _interrupt_signal_handle_callback(self, signal, frame): """Handle the interrupt signal from the keyboard.""" diff --git a/source/isaaclab/isaaclab/assets/articulation/base_articulation.py b/source/isaaclab/isaaclab/assets/articulation/base_articulation.py index 25ca2c4ceaf0..14aa592ad103 100644 --- a/source/isaaclab/isaaclab/assets/articulation/base_articulation.py +++ b/source/isaaclab/isaaclab/assets/articulation/base_articulation.py @@ -2553,14 +2553,17 @@ def set_external_force_and_torque( env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, is_global: bool = False, ) -> None: - """Deprecated, same as :meth:`permanent_wrench_composer.set_forces_and_torques`.""" + """Deprecated. Resets target environments, then adds forces and torques via the permanent wrench composer.""" warnings.warn( - "The function 'set_external_force_and_torque' will be deprecated in a future release. Please" - " use 'permanent_wrench_composer.set_forces_and_torques' instead.", + "The function 'set_external_force_and_torque' is deprecated. Please use" + " 'permanent_wrench_composer.reset' followed by 'permanent_wrench_composer.add_forces_and_torques'" + " instead.", DeprecationWarning, stacklevel=2, ) - self.permanent_wrench_composer.set_forces_and_torques( + # Reset only target env_ids then add (not set which clears all envs globally) + self.permanent_wrench_composer.reset(env_ids=env_ids) + self.permanent_wrench_composer.add_forces_and_torques( forces, torques, positions=positions, body_ids=body_ids, env_ids=env_ids, is_global=is_global ) diff --git a/source/isaaclab/isaaclab/assets/rigid_object/base_rigid_object.py b/source/isaaclab/isaaclab/assets/rigid_object/base_rigid_object.py index a2e57aed5419..d7bc6ea4c85d 100644 --- a/source/isaaclab/isaaclab/assets/rigid_object/base_rigid_object.py +++ b/source/isaaclab/isaaclab/assets/rigid_object/base_rigid_object.py @@ -845,13 +845,16 @@ def set_external_force_and_torque( env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, is_global: bool = False, ) -> None: - """Deprecated, same as :meth:`permanent_wrench_composer.set_forces_and_torques`.""" + """Deprecated. Resets target environments, then adds forces and torques via the permanent wrench composer.""" warnings.warn( - "The function 'set_external_force_and_torque' will be deprecated in a future release. Please" - " use 'permanent_wrench_composer.set_forces_and_torques' instead.", + "The function 'set_external_force_and_torque' is deprecated. Please use" + " 'permanent_wrench_composer.reset' followed by 'permanent_wrench_composer.add_forces_and_torques'" + " instead.", DeprecationWarning, stacklevel=2, ) - self.permanent_wrench_composer.set_forces_and_torques( + # Reset only target env_ids then add (not set which clears all envs globally) + self.permanent_wrench_composer.reset(env_ids=env_ids) + self.permanent_wrench_composer.add_forces_and_torques( forces, torques, positions=positions, body_ids=body_ids, env_ids=env_ids, is_global=is_global ) diff --git a/source/isaaclab/isaaclab/assets/rigid_object_collection/base_rigid_object_collection.py b/source/isaaclab/isaaclab/assets/rigid_object_collection/base_rigid_object_collection.py index 0384bc67da58..4be586145b0e 100644 --- a/source/isaaclab/isaaclab/assets/rigid_object_collection/base_rigid_object_collection.py +++ b/source/isaaclab/isaaclab/assets/rigid_object_collection/base_rigid_object_collection.py @@ -904,20 +904,18 @@ def set_external_force_and_torque( env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, is_global: bool = False, ) -> None: - """Deprecated, same as :meth:`permanent_wrench_composer.set_forces_and_torques`.""" + """Deprecated. Resets target environments, then adds forces and torques via the permanent wrench composer.""" warnings.warn( - "The function 'set_external_force_and_torque' will be deprecated in a future release. Please" - " use 'permanent_wrench_composer.set_forces_and_torques' instead.", + "The function 'set_external_force_and_torque' is deprecated. Please use" + " 'permanent_wrench_composer.reset' followed by 'permanent_wrench_composer.add_forces_and_torques'" + " instead.", DeprecationWarning, stacklevel=2, ) - self.permanent_wrench_composer.set_forces_and_torques( - forces=forces, - torques=torques, - positions=positions, - body_ids=body_ids, - env_ids=env_ids, - is_global=is_global, + # Reset only target env_ids then add (not set which clears all envs globally) + self.permanent_wrench_composer.reset(env_ids=env_ids) + self.permanent_wrench_composer.add_forces_and_torques( + forces, torques, positions=positions, body_ids=body_ids, env_ids=env_ids, is_global=is_global ) def write_object_state_to_sim( diff --git a/source/isaaclab/isaaclab/cli/commands/install.py b/source/isaaclab/isaaclab/cli/commands/install.py index c442cfd89fe6..53cf4a799fb1 100644 --- a/source/isaaclab/isaaclab/cli/commands/install.py +++ b/source/isaaclab/isaaclab/cli/commands/install.py @@ -305,6 +305,9 @@ def _install_extra_frameworks(framework_name: str = "all") -> None: "newton_actuators", "warp", "mujoco_warp", + "websockets", + "viser", + "imgui_bundle", ] """Package directory names in Isaac Sim prebundle directories to repoint. @@ -352,7 +355,25 @@ def _repoint_prebundle_packages() -> None: print_warning(f"site-packages directory not found: {site_packages} — skipping prebundle repoint.") return - prebundle_dirs = list(isaacsim_path.rglob("pip_prebundle")) + # Discover pip_prebundle directories from both the Isaac Sim tree and + # Omniverse cache roots. Some Isaac Sim directories are symlinked into + # ~/.local/share/ov and may be missed by a plain rglob() on _isaac_sim. + candidate_roots: set[Path] = set() + for root in ( + isaacsim_path, + isaacsim_path.resolve(), + isaacsim_path / "extscache", + Path.home() / ".local" / "share" / "ov" / "data" / "exts", + Path.home() / ".local" / "share" / "ov" / "data" / "exts" / "v2", + ): + if root.exists(): + candidate_roots.add(root) + candidate_roots.add(root.resolve()) + + prebundle_dirs: set[Path] = set() + for root in candidate_roots: + prebundle_dirs.update(root.rglob("pip_prebundle")) + if not prebundle_dirs: print_debug("No pip_prebundle directories found under Isaac Sim.") return @@ -482,12 +503,6 @@ def command_install(install_type: str = "all") -> None: if name == "newton" and "isaaclab_visualizers" not in isaaclab_submodules: isaaclab_submodules.append("isaaclab_visualizers") submodule_extras["isaaclab_visualizers"] = "[newton]" - # newton and physx are tightly coupled; always install both together. - # todo: remove once we move to UV and pyproject.toml-based packaging - if name == "newton" and "isaaclab_physx" not in isaaclab_submodules: - isaaclab_submodules.append("isaaclab_physx") - if name == "physx" and "isaaclab_newton" not in isaaclab_submodules: - isaaclab_submodules.append("isaaclab_newton") else: valid = sorted(VALID_ISAACLAB_SUBMODULES) + sorted(VALID_RL_FRAMEWORKS) + ["isaacsim"] print_warning(f"Unknown Isaac Lab submodule '{name}'. Valid values: {', '.join(valid)}. Skipping.") diff --git a/source/isaaclab/isaaclab/cli/utils.py b/source/isaaclab/isaaclab/cli/utils.py index a1870b9c2225..5a3b60532870 100644 --- a/source/isaaclab/isaaclab/cli/utils.py +++ b/source/isaaclab/isaaclab/cli/utils.py @@ -146,19 +146,16 @@ def _print_debug_env(prefix: str, env: dict[str, str] | None) -> None: _CMD_METACHARACTERS = frozenset("<>|&^") -def _escape_for_cmd_exe(cmd: list[str] | tuple[str, ...]) -> str | list[str]: +def _escape_for_cmd_exe(cmd: list[str] | tuple[str, ...]) -> list[str]: + """Wrap .bat/.cmd calls in cmd.exe /c so args with < > | & ^ stay literal + (otherwise Windows treats e.g. setuptools<82.0.0 as a redirection). """ - Quote ``cmd.exe`` metacharacters when invoking ``.bat``/``.cmd`` files. - - Returns a command string (not list) so ``subprocess.run`` bypasses - ``list2cmdline`` which doesn't escape cmd.exe metacharacters (<, >, |, &, ^). - """ - # Only .bat/.cmd files are executed via cmd.exe. + # only .bat/.cmd needs wrapping exe = str(cmd[0]).lower() if not (exe.endswith(".bat") or exe.endswith(".cmd")): return list(cmd) - # Wrap args that contain metacharacters or whitespace in double quotes. + # quote anything with metacharacters or spaces parts: list[str] = [] for arg in cmd: s = str(arg) @@ -167,8 +164,7 @@ def _escape_for_cmd_exe(cmd: list[str] | tuple[str, ...]) -> str | list[str]: else: parts.append(s) - # Return a string so subprocess skips list2cmdline. - return " ".join(parts) + return ["cmd.exe", "/c", " ".join(parts)] def run_command( diff --git a/source/isaaclab/isaaclab/envs/common.py b/source/isaaclab/isaaclab/envs/common.py index f913005d1dbb..5da6f871361e 100644 --- a/source/isaaclab/isaaclab/envs/common.py +++ b/source/isaaclab/isaaclab/envs/common.py @@ -5,6 +5,8 @@ from __future__ import annotations +import warnings +from dataclasses import MISSING, fields from typing import Dict, Literal, TypeVar # noqa: UP035 import gymnasium as gym @@ -17,9 +19,25 @@ ## +def _viewer_cfg_value_matches_default(current: object, default: object) -> bool: + """Return True if ``current`` matches the dataclass field default (including list/tuple equivalence).""" + if current == default: + return True + if isinstance(current, (list, tuple)) and isinstance(default, (list, tuple)): + if len(current) != len(default): + return False + return all(a == b for a, b in zip(current, default, strict=True)) + return False + + @configclass class ViewerCfg: - """Configuration of the scene viewport camera.""" + """Configuration of the scene viewport camera. + + Note: + ViewerCfg is deprecated. In a future release, this config will be streamlined with + the KitVisualizerCfg. + """ eye: tuple[float, float, float] = (7.5, 7.5, 7.5) """Initial camera position (in m). Default is (7.5, 7.5, 7.5).""" @@ -67,6 +85,31 @@ class ViewerCfg: This quantity is only effective if :attr:`origin` is set to "asset_body". """ + def __post_init__(self) -> None: + # Dataclasses do not record which arguments were passed explicitly vs defaulted, and + # warning only on ``**kwargs`` would miss positional arguments. Comparing each field to + # its declared default catches any non-default effective configuration (including + # ``replace()`` and ``from_dict``), while keeping ``ViewerCfg()`` silent. + differing: list[str] = [] + for f in fields(self): + if not f.init: + continue + if f.default is not MISSING: + default_val = f.default + elif f.default_factory is not MISSING: + default_val = f.default_factory() + else: + continue + if not _viewer_cfg_value_matches_default(getattr(self, f.name), default_val): + differing.append(f.name) + if differing: + warnings.warn( + "ViewerCfg is deprecated. In a future release, this config will be streamlined with " + "the KitVisualizerCfg.", + DeprecationWarning, + stacklevel=2, + ) + ## # Types. diff --git a/source/isaaclab/isaaclab/envs/direct_marl_env.py b/source/isaaclab/isaaclab/envs/direct_marl_env.py index c6009117f1b5..b325164ebe01 100644 --- a/source/isaaclab/isaaclab/envs/direct_marl_env.py +++ b/source/isaaclab/isaaclab/envs/direct_marl_env.py @@ -151,10 +151,10 @@ def _init_sim(self, render_mode: str | None = None, **kwargs): # viewport is not available in other rendering modes so the function will throw a warning # FIXME: This needs to be fixed in the future when we unify the UI functionalities even for # non-rendering modes. - # Initialize when GUI is available OR when visualizers are active (headless rendering) - # Visualizers support camera updates via sim.set_camera_view() which forwards to all active visualizers - has_visualizers = bool(self.sim.get_setting("/isaaclab/visualizer")) - if self.sim.has_gui or has_visualizers: + # Initialize when a Kit viewport exists. ViewportCameraController uses omni.kit (renderer camera); + # skip in kitless Newton-only runs (e.g. --viz rerun) where no Kit app is running. + has_visualizers = self.sim.has_active_visualizers() + if (self.sim.has_gui or has_visualizers) and has_kit(): self.viewport_camera_controller = ViewportCameraController(self, self.cfg.viewer) else: self.viewport_camera_controller = None diff --git a/source/isaaclab/isaaclab/envs/direct_rl_env.py b/source/isaaclab/isaaclab/envs/direct_rl_env.py index 05dc8495dbcf..c67803ff8cf1 100644 --- a/source/isaaclab/isaaclab/envs/direct_rl_env.py +++ b/source/isaaclab/isaaclab/envs/direct_rl_env.py @@ -156,10 +156,10 @@ def _init_sim(self, render_mode: str | None = None, **kwargs): # viewport is not available in other rendering modes so the function will throw a warning # FIXME: This needs to be fixed in the future when we unify the UI functionalities even for # non-rendering modes. - # Initialize when GUI is available OR when visualizers are active (headless rendering) - # Visualizers support camera updates via sim.set_camera_view() which forwards to all active visualizers - has_visualizers = bool(self.sim.get_setting("/isaaclab/visualizer")) - if self.sim.has_gui or has_visualizers: + # Initialize when a Kit viewport exists. ViewportCameraController uses omni.kit (renderer camera); + # skip in kitless Newton-only runs (e.g. --viz rerun) where no Kit app is running. + has_visualizers = self.sim.has_active_visualizers() + if (self.sim.has_gui or has_visualizers) and has_kit(): self.viewport_camera_controller = ViewportCameraController(self, self.cfg.viewer) else: self.viewport_camera_controller = None diff --git a/source/isaaclab/isaaclab/envs/manager_based_env.py b/source/isaaclab/isaaclab/envs/manager_based_env.py index 1e8ca0576101..c63db4922c9e 100644 --- a/source/isaaclab/isaaclab/envs/manager_based_env.py +++ b/source/isaaclab/isaaclab/envs/manager_based_env.py @@ -21,6 +21,7 @@ from isaaclab.utils.configclass import resolve_cfg_presets from isaaclab.utils.seed import configure_seed from isaaclab.utils.timer import Timer +from isaaclab.utils.version import has_kit from .common import VecEnvObs from .manager_based_env_cfg import ManagerBasedEnvCfg @@ -166,10 +167,10 @@ def _init_sim(self): # viewport is not available in other rendering modes so the function will throw a warning # FIXME: This needs to be fixed in the future when we unify the UI functionalities even for # non-rendering modes. - # Initialize when GUI is available OR when visualizers are active (headless rendering) - # Visualizers support camera updates via sim.set_camera_view() which forwards to all active visualizers - has_visualizers = bool(self.sim.get_setting("/isaaclab/visualizer")) - if self.sim.has_gui or has_visualizers: + # Initialize when a Kit viewport exists. ViewportCameraController uses omni.kit (renderer camera); + # skip in kitless Newton-only runs (e.g. --viz rerun) where no Kit app is running. + has_visualizers = self.sim.has_active_visualizers() + if (self.sim.has_gui or has_visualizers) and has_kit(): self.viewport_camera_controller = ViewportCameraController(self, self.cfg.viewer) else: self.viewport_camera_controller = None diff --git a/source/isaaclab/isaaclab/envs/mdp/observations.py b/source/isaaclab/isaaclab/envs/mdp/observations.py index a207749550a9..d65bd264f7ed 100644 --- a/source/isaaclab/isaaclab/envs/mdp/observations.py +++ b/source/isaaclab/isaaclab/envs/mdp/observations.py @@ -304,7 +304,7 @@ def height_scan(env: ManagerBasedEnv, sensor_cfg: SceneEntityCfg, offset: float # extract the used quantities (to enable type-hinting) sensor: RayCaster = env.scene.sensors[sensor_cfg.name] # height scan: height = sensor_height - hit_point_z - offset - return sensor.data.pos_w[:, 2].unsqueeze(1) - sensor.data.ray_hits_w[..., 2] - offset + return wp.to_torch(sensor.data.pos_w)[:, 2].unsqueeze(1) - wp.to_torch(sensor.data.ray_hits_w)[..., 2] - offset def body_incoming_wrench(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg) -> torch.Tensor: diff --git a/source/isaaclab/isaaclab/envs/mdp/rewards.py b/source/isaaclab/isaaclab/envs/mdp/rewards.py index 5a53583f5e4b..74bea7ee7861 100644 --- a/source/isaaclab/isaaclab/envs/mdp/rewards.py +++ b/source/isaaclab/isaaclab/envs/mdp/rewards.py @@ -116,7 +116,7 @@ def base_height_l2( if sensor_cfg is not None: sensor: RayCaster = env.scene[sensor_cfg.name] # Adjust the target height using the sensor data - adjusted_target_height = target_height + torch.mean(sensor.data.ray_hits_w[..., 2], dim=1) + adjusted_target_height = target_height + torch.mean(wp.to_torch(sensor.data.ray_hits_w)[..., 2], dim=1) else: # Use the provided target height directly for flat terrain adjusted_target_height = target_height diff --git a/source/isaaclab/isaaclab/envs/ui/viewport_camera_controller.py b/source/isaaclab/isaaclab/envs/ui/viewport_camera_controller.py index 4126d7b74735..277982e7a2c9 100644 --- a/source/isaaclab/isaaclab/envs/ui/viewport_camera_controller.py +++ b/source/isaaclab/isaaclab/envs/ui/viewport_camera_controller.py @@ -218,8 +218,17 @@ def update_view_location(self, eye: Sequence[float] | None = None, lookat: Seque cam_eye = viewer_origin + self.default_cam_eye cam_target = viewer_origin + self.default_cam_lookat - # set the camera view - self._env.sim.set_camera_view(eye=cam_eye, target=cam_target) + eye_t = (float(cam_eye[0]), float(cam_eye[1]), float(cam_eye[2])) + target_t = (float(cam_target[0]), float(cam_target[1]), float(cam_target[2])) + self._env.sim.set_camera_view(eye=eye_t, target=target_t) + + # Renderer viewport camera (Isaac RTX / Kit); optional — pure-Newton installs have no isaaclab_physx. + try: + from isaaclab_physx.renderers.kit_viewport_utils import set_kit_renderer_camera_view + + set_kit_renderer_camera_view(eye=cam_eye, target=cam_target, camera_prim_path=self.cfg.cam_prim_path) + except (ImportError, ModuleNotFoundError): + pass """ Private Functions diff --git a/source/isaaclab/isaaclab/envs/utils/recording_hooks.py b/source/isaaclab/isaaclab/envs/utils/recording_hooks.py new file mode 100644 index 000000000000..584b6d73adf5 --- /dev/null +++ b/source/isaaclab/isaaclab/envs/utils/recording_hooks.py @@ -0,0 +1,50 @@ +# 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 + +"""Hooks that run after visualizers during :meth:`~isaaclab.sim.SimulationContext.render`. + +Lives alongside :mod:`video_recorder` / :mod:`video_recorder_cfg` because both tie into +``--video`` / ``rgb_array`` recording. Keeps :class:`~isaaclab.sim.SimulationContext` free +of imports from ``isaaclab_physx``, ``isaaclab_newton``, and other recording backends. +Each integration is loaded lazily so optional extensions are not required at import time. +""" + +from __future__ import annotations + +from typing import Any + + +def run_recording_hooks_after_visualizers(sim: Any) -> None: + """Run recording-related work after :meth:`~isaaclab.sim.SimulationContext.render` steps visualizers. + + Isaac Sim / RTX follow-up is loaded lazily so minimal installs still work. + Newton GL video is handled by :class:`~isaaclab.envs.utils.video_recorder.VideoRecorder` + (e.g. :class:`~isaaclab_newton.video_recording.newton_gl_perspective_video.NewtonGlPerspectiveVideo`), + not here. + + Args: + sim: Active :class:`~isaaclab.sim.SimulationContext` instance. + """ + _recording_followup_isaac_sim(sim) + + +def _recording_followup_isaac_sim(sim: Any) -> None: + """Isaac Sim: keep RTX / Replicator outputs fresh when recording video without a Kit visualizer. + + When ``--video`` uses ``rgb_array`` / :class:`~gymnasium.wrappers.RecordVideo`, Replicator + render products must see Kit's event loop pumped. :class:`~isaaclab_visualizers.kit.KitVisualizer` + already calls ``omni.kit.app.get_app().update()`` in its ``step()``; if no such visualizer + is active, we pump here (guarded by ``/isaaclab/video/enabled`` and ``is_rendering``). + + Implemented by ``pump_kit_app_for_headless_video_render_if_needed`` in + :mod:`isaaclab_physx.renderers.isaac_rtx_renderer_utils`. + """ + try: + from isaaclab_physx.renderers.isaac_rtx_renderer_utils import ( + pump_kit_app_for_headless_video_render_if_needed, + ) + except ImportError: + return + pump_kit_app_for_headless_video_render_if_needed(sim) diff --git a/source/isaaclab/isaaclab/physics/base_scene_data_provider.py b/source/isaaclab/isaaclab/physics/base_scene_data_provider.py index e5b709da0ce7..9760a71d25ba 100644 --- a/source/isaaclab/isaaclab/physics/base_scene_data_provider.py +++ b/source/isaaclab/isaaclab/physics/base_scene_data_provider.py @@ -15,8 +15,8 @@ class BaseSceneDataProvider(ABC): """Backend-agnostic scene data provider interface.""" @abstractmethod - def update(self, env_ids: list[int] | None = None) -> None: - """Refresh any cached scene data.""" + def update(self) -> None: + """Refresh any cached scene data (full model/state).""" raise NotImplementedError @abstractmethod @@ -25,8 +25,8 @@ def get_newton_model(self) -> Any | None: raise NotImplementedError @abstractmethod - def get_newton_state(self, env_ids: list[int] | None = None) -> Any | None: - """Return Newton state handle when available.""" + def get_newton_state(self) -> Any | None: + """Return Newton state handle when available (full state).""" raise NotImplementedError @abstractmethod diff --git a/source/isaaclab/isaaclab/physics/physics_manager.py b/source/isaaclab/isaaclab/physics/physics_manager.py index cc18582bc80e..7a4cdfe84403 100644 --- a/source/isaaclab/isaaclab/physics/physics_manager.py +++ b/source/isaaclab/isaaclab/physics/physics_manager.py @@ -276,6 +276,16 @@ def pre_render(cls) -> None: """ pass + @classmethod + def after_visualizers_render(cls) -> None: + """Hook after visualizers have stepped during :meth:`~isaaclab.sim.SimulationContext.render`. + + Use for physics-backend sync (e.g. fabric) if needed. Recording pipelines (Kit/RTX, + Newton GL video, etc.) run from :mod:`isaaclab.envs.utils.recording_hooks` so they are not + tied to a specific physics manager. Default is a no-op. + """ + pass + @classmethod def close(cls) -> None: """Clean up physics resources. 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/kernels.py b/source/isaaclab/isaaclab/sensors/ray_caster/kernels.py new file mode 100644 index 000000000000..98c54ea7141d --- /dev/null +++ b/source/isaaclab/isaaclab/sensors/ray_caster/kernels.py @@ -0,0 +1,261 @@ +# 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 + +"""Warp kernels for the ray caster sensor.""" + +import warp as wp + +ALIGNMENT_WORLD = wp.constant(0) +ALIGNMENT_YAW = wp.constant(1) +ALIGNMENT_BASE = wp.constant(2) + +# Upper-bound ray-cast distance [m] used by camera classes. The actual depth-clipping is applied +# as a post-process step per data type, so the kernel is always given a large budget. +CAMERA_RAYCAST_MAX_DIST: float = 1e6 + + +@wp.func +def quat_yaw_only( + # input + q: wp.quatf, +) -> wp.quatf: + """Extract the yaw-only quaternion from a general quaternion. + + Equivalent to :func:`isaaclab.utils.math.yaw_quat`: extracts the yaw angle via + ``atan2(2*(qw*qz + qx*qy), 1 - 2*(qy^2 + qz^2))`` and returns a pure-yaw quaternion + ``(0, 0, sin(yaw/2), cos(yaw/2))``. This is correct for all orientations, including + those with non-zero roll and pitch. + """ + qx = q[0] + qy = q[1] + qz = q[2] + qw = q[3] + yaw = wp.atan2(2.0 * (qw * qz + qx * qy), 1.0 - 2.0 * (qy * qy + qz * qz)) + half_yaw = yaw * 0.5 + return wp.quatf(0.0, 0.0, wp.sin(half_yaw), wp.cos(half_yaw)) + + +@wp.kernel(enable_backward=False) +def update_ray_caster_kernel( + # input + transforms: wp.array(dtype=wp.transformf), + env_mask: wp.array(dtype=wp.bool), + offset_pos: wp.array(dtype=wp.vec3f), + offset_quat: wp.array(dtype=wp.quatf), + drift: wp.array(dtype=wp.vec3f), + ray_cast_drift: wp.array(dtype=wp.vec3f), + ray_starts_local: wp.array2d(dtype=wp.vec3f), + ray_directions_local: wp.array2d(dtype=wp.vec3f), + alignment_mode: int, + # output + pos_w: wp.array(dtype=wp.vec3f), + quat_w: wp.array(dtype=wp.quatf), + ray_starts_w: wp.array2d(dtype=wp.vec3f), + ray_directions_w: wp.array2d(dtype=wp.vec3f), +): + """Compute sensor world poses and transform rays into world frame. + + Combines the PhysX view transform with the sensor offset, applies drift, + and transforms local ray starts/directions according to the alignment mode. + + Launch with dim=(num_envs, num_rays). + + Args: + transforms: World transforms from PhysX view. Shape is (num_envs,). + env_mask: Boolean mask for which environments to update. Shape is (num_envs,). + offset_pos: Per-env position offset [m] from view to sensor. Shape is (num_envs,). + offset_quat: Per-env quaternion offset from view to sensor. Shape is (num_envs,). + drift: Per-env position drift [m]. Shape is (num_envs,). + ray_cast_drift: Per-env ray cast drift [m]. Shape is (num_envs,). + After rotation by the alignment quaternion, only the x and y components + are applied to the ray start position; the z component of the sensor + position is preserved. + ray_starts_local: Per-env local ray start positions [m]. Shape is (num_envs, num_rays). + ray_directions_local: Per-env local ray directions (unit vectors). Shape is (num_envs, num_rays). + alignment_mode: 0=world, 1=yaw, 2=base. + pos_w: Output sensor position in world frame [m]. Shape is (num_envs,). + quat_w: Output sensor orientation in world frame. Shape is (num_envs,). + ray_starts_w: Output world-frame ray starts [m]. Shape is (num_envs, num_rays). + ray_directions_w: Output world-frame ray directions (unit vectors). Shape is (num_envs, num_rays). + """ + env_id, ray_id = wp.tid() + if not env_mask[env_id]: + return + + t = transforms[env_id] + view_pos = wp.transform_get_translation(t) + view_quat = wp.transform_get_rotation(t) + + # combine_frame_transforms: q02 = q01 * q12, t02 = t01 + quat_rotate(q01, t12) + combined_quat = view_quat * offset_quat[env_id] + combined_pos = view_pos + wp.quat_rotate(view_quat, offset_pos[env_id]) + + combined_pos = combined_pos + drift[env_id] + + if ray_id == 0: + pos_w[env_id] = combined_pos + quat_w[env_id] = combined_quat + + local_start = ray_starts_local[env_id, ray_id] + local_dir = ray_directions_local[env_id, ray_id] + rcd = ray_cast_drift[env_id] + + if alignment_mode == ALIGNMENT_WORLD: + pos_drifted = wp.vec3f(combined_pos[0] + rcd[0], combined_pos[1] + rcd[1], combined_pos[2]) + ray_starts_w[env_id, ray_id] = local_start + pos_drifted + ray_directions_w[env_id, ray_id] = local_dir + elif alignment_mode == ALIGNMENT_YAW: + yaw_q = quat_yaw_only(combined_quat) + rot_drift = wp.quat_rotate(yaw_q, rcd) + pos_drifted = wp.vec3f(combined_pos[0] + rot_drift[0], combined_pos[1] + rot_drift[1], combined_pos[2]) + ray_starts_w[env_id, ray_id] = wp.quat_rotate(yaw_q, local_start) + pos_drifted + # Ray DIRECTIONS are intentionally NOT rotated in yaw mode: the sensor's ray pattern + # (e.g. straight-down (0,0,-1) for a height scanner) stays fixed in world frame. + # Only ray STARTS are rotated by the yaw-only quaternion so the scan footprint + # follows the body heading without tilting when the body pitches or rolls. + ray_directions_w[env_id, ray_id] = local_dir + else: + rot_drift = wp.quat_rotate(combined_quat, rcd) + pos_drifted = wp.vec3f(combined_pos[0] + rot_drift[0], combined_pos[1] + rot_drift[1], combined_pos[2]) + ray_starts_w[env_id, ray_id] = wp.quat_rotate(combined_quat, local_start) + pos_drifted + ray_directions_w[env_id, ray_id] = wp.quat_rotate(combined_quat, local_dir) + + +@wp.kernel(enable_backward=False) +def fill_vec3_inf_kernel( + # input + env_mask: wp.array(dtype=wp.bool), + inf_val: wp.float32, + # output + data: wp.array2d(dtype=wp.vec3f), +): + """Fill a 2D vec3f array with a given value for masked environments. + + Launch with dim=(num_envs, num_rays). + + Args: + env_mask: Boolean mask for which environments to update. Shape is (num_envs,). + inf_val: Value to fill with (typically inf). + data: Array to fill. Shape is (num_envs, num_rays). + """ + env, ray = wp.tid() + if not env_mask[env]: + return + data[env, ray] = wp.vec3f(inf_val, inf_val, inf_val) + + +@wp.kernel(enable_backward=False) +def apply_z_drift_kernel( + # input + env_mask: wp.array(dtype=wp.bool), + ray_cast_drift: wp.array(dtype=wp.vec3f), + # output + ray_hits: wp.array2d(dtype=wp.vec3f), +): + """Apply vertical (z) drift to ray hit positions for masked environments. + + Launch with dim=(num_envs, num_rays). + + Args: + env_mask: Boolean mask for which environments to update. Shape is (num_envs,). + ray_cast_drift: Per-env drift vector [m]; only z-component is used. Shape is (num_envs,). + ray_hits: Ray hit positions to modify in-place. Shape is (num_envs, num_rays). + """ + env, ray = wp.tid() + if not env_mask[env]: + return + hit = ray_hits[env, ray] + ray_hits[env, ray] = wp.vec3f(hit[0], hit[1], hit[2] + ray_cast_drift[env][2]) + + +@wp.kernel(enable_backward=False) +def fill_float2d_masked_kernel( + # input + env_mask: wp.array(dtype=wp.bool), + val: wp.float32, + # output + data: wp.array2d(dtype=wp.float32), +): + """Fill a 2D float32 array with a given value for masked environments. + + Launch with dim=(num_envs, num_rays). + + Args: + env_mask: Boolean mask for which environments to update. Shape is (num_envs,). + val: Value to fill with. + data: Array to fill. Shape is (num_envs, num_rays). + """ + env, ray = wp.tid() + if not env_mask[env]: + return + data[env, ray] = val + + +@wp.kernel(enable_backward=False) +def compute_distance_to_image_plane_masked_kernel( + # input + env_mask: wp.array(dtype=wp.bool), + quat_w: wp.array(dtype=wp.quatf), + ray_distance: wp.array2d(dtype=wp.float32), + ray_directions_w: wp.array2d(dtype=wp.vec3f), + # output + distance_to_image_plane: wp.array2d(dtype=wp.float32), +): + """Compute distance-to-image-plane from ray depth and direction for masked environments. + + The distance to the image plane is the signed projection of the hit displacement + (``ray_distance * ray_direction_w``) onto the camera forward axis (+X in world convention). + This equals the x-component of the hit vector in the camera frame. + + Launch with dim=(num_envs, num_rays). + + Args: + env_mask: Boolean mask for which environments to update. Shape is (num_envs,). + quat_w: Camera orientation in world frame (x, y, z, w). Shape is (num_envs,). + ray_distance: Per-ray hit distances [m]. Shape is (num_envs, num_rays). + Contains inf for missed rays. + ray_directions_w: World-frame unit ray directions. Shape is (num_envs, num_rays). + distance_to_image_plane: Output distance-to-image-plane [m]. Shape is (num_envs, num_rays). + """ + env, ray = wp.tid() + if not env_mask[env]: + return + + depth = ray_distance[env, ray] + dir_w = ray_directions_w[env, ray] + # displacement vector in world frame + disp_w = wp.vec3f(depth * dir_w[0], depth * dir_w[1], depth * dir_w[2]) + # rotate into camera frame (quat_rotate_inv applies q^-1 * v * q) + disp_cam = wp.quat_rotate_inv(quat_w[env], disp_w) + # x-component is the forward (depth) axis of the camera in world convention + distance_to_image_plane[env, ray] = disp_cam[0] + + +@wp.kernel(enable_backward=False) +def apply_depth_clipping_masked_kernel( + # input + env_mask: wp.array(dtype=wp.bool), + max_dist: wp.float32, + fill_val: wp.float32, + # output + depth: wp.array2d(dtype=wp.float32), +): + """Clip depth values in-place, replacing values above max_dist or NaN with fill_val. + + Launch with dim=(num_envs, num_rays). + + Args: + env_mask: Boolean mask for which environments to update. Shape is (num_envs,). + max_dist: Maximum depth threshold [m]. + fill_val: Replacement value [m] written for depths exceeding max_dist or NaN. + Pass ``max_dist`` for "max" clipping or ``0.0`` for "zero" clipping. + depth: Depth buffer to clip in-place. Shape is (num_envs, num_rays). + """ + env, ray = wp.tid() + if not env_mask[env]: + return + val = depth[env, ray] + if val > max_dist or wp.isnan(val): + depth[env, ray] = fill_val 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 06ce2183e2ff..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,22 +14,20 @@ 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 matrix_from_quat, quat_mul +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, raycast_dynamic_meshes +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: from .multi_mesh_ray_caster_cfg import MultiMeshRayCasterCfg -# import logger logger = logging.getLogger(__name__) @@ -41,8 +39,8 @@ class MultiMeshRayCaster(RayCaster): a set of meshes with a given ray pattern. The meshes are parsed from the list of primitive paths provided in the configuration. These are then - converted to warp meshes and stored in the :attr:`meshes` list. The ray-caster then ray-casts against - these warp meshes using the ray pattern provided in the configuration. + converted to warp meshes and stored in the :attr:`meshes` dictionary. The ray-caster then ray-casts + against these warp meshes using the ray pattern provided in the configuration. Compared to the default RayCaster, the MultiMeshRayCaster provides additional functionality and flexibility as an extension of the default RayCaster with the following enhancements: @@ -53,6 +51,15 @@ class MultiMeshRayCaster(RayCaster): (e.g., robot links, articulated bodies, or dynamic obstacles). - Memory-efficient caching : Avoids redundant memory usage by reusing mesh data across environments. + .. warning:: + **Known limitation (multi-mesh closest-hit resolution):** When two meshes produce a + hit at the exact same distance for a given ray, the ``atomic_min`` + equality-check + pattern in the raycasting kernel is not fully thread-safe. The hit *position* is always + correct, but auxiliary outputs (normals, face IDs, mesh IDs) may originate from + different meshes for the affected ray. This requires an exact floating-point tie and is + rare in practice. See `warp#1058 `_ for + upstream progress on a thread-safe ``atomic_min`` return value. + Example usage to raycast against the visual meshes of a robot (e.g. ANYmal): .. code-block:: python @@ -76,9 +83,7 @@ class MultiMeshRayCaster(RayCaster): cfg: MultiMeshRayCasterCfg """The configuration parameters.""" - mesh_offsets: dict[str, tuple[torch.Tensor, torch.Tensor]] = {} - - 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. @@ -90,33 +95,24 @@ def __init__(self, cfg: MultiMeshRayCasterCfg): Args: cfg: The configuration parameters. """ - # Initialize base class super().__init__(cfg) - # Create empty variables for storing output data self._num_meshes_per_env: dict[str, int] = {} - """Keeps track of the number of meshes per env for each ray_cast target. - Since we allow regex indexing (e.g. env_*/object_*) they can differ - """ self._raycast_targets_cfg: list[MultiMeshRayCasterCfg.RaycastTargetCfg] = [] for target in self.cfg.mesh_prim_paths: - # Legacy support for string targets. Treat them as global targets. if isinstance(target, str): self._raycast_targets_cfg.append(cfg.RaycastTargetCfg(prim_expr=target, track_mesh_transforms=False)) else: self._raycast_targets_cfg.append(target) - # Resolve regex namespace if set for cfg in self._raycast_targets_cfg: cfg.prim_expr = cfg.prim_expr.format(ENV_REGEX_NS="/World/envs/env_.*") - # overwrite the data class self._data = MultiMeshRayCasterData() def __str__(self) -> str: """Returns: A string containing information about the instance.""" - return ( f"Ray-caster @ '{self.cfg.prim_path}': \n" f"\tview type : {self._view.__class__}\n" @@ -133,9 +129,7 @@ def __str__(self) -> str: @property def data(self) -> MultiMeshRayCasterData: - # update sensors if needed self._update_outdated_buffers() - # return the data return self._data """ @@ -163,9 +157,8 @@ def _initialize_warp_meshes(self): """ multi_mesh_ids: dict[str, list[list[int]]] = {} for target_cfg in self._raycast_targets_cfg: - # target prim path to ray cast against target_prim_path = target_cfg.prim_expr - # # check if mesh already casted into warp mesh and skip if so. + # check if mesh already casted into warp mesh and skip if so. if target_prim_path in multi_mesh_ids: logger.warning( f"Mesh at target prim path '{target_prim_path}' already exists in the mesh cache. Duplicate entries" @@ -173,32 +166,29 @@ def _initialize_warp_meshes(self): ) continue - # find all matching prim paths to provided expression of the target target_prims = sim_utils.find_matching_prims(target_prim_path) if len(target_prims) == 0: raise RuntimeError(f"Failed to find a prim at path expression: {target_prim_path}") - # If only one prim is found, treat it as a global prim. - # Either it's a single global object (e.g. ground) or we are only using one env. is_global_prim = len(target_prims) == 1 loaded_vertices: list[np.ndarray | None] = [] wp_mesh_ids = [] for target_prim in target_prims: - # Reuse previously parsed shared mesh instance if possible. if target_cfg.is_shared and len(wp_mesh_ids) > 0: # Verify if this mesh has already been registered in an earlier environment. # Note, this check may fail, if the prim path is not following the env_.* pattern # Which (worst case) leads to parsing the mesh and skipping registering it at a later stage - curr_prim_base_path = re.sub(r"env_\d+", "env_0", str(target_prim.GetPath())) # - if curr_prim_base_path in MultiMeshRayCaster.meshes: - MultiMeshRayCaster.meshes[str(target_prim.GetPath())] = MultiMeshRayCaster.meshes[ - curr_prim_base_path - ] - # Reuse mesh imported by another ray-cast sensor (global cache). - if str(target_prim.GetPath()) in MultiMeshRayCaster.meshes: - wp_mesh_ids.append(MultiMeshRayCaster.meshes[str(target_prim.GetPath())].id) + curr_prim_base_path = re.sub(r"env_\d+", "env_0", str(target_prim.GetPath())) + base_key = (curr_prim_base_path, self._device) + if base_key in MultiMeshRayCaster.meshes: + MultiMeshRayCaster.meshes[(str(target_prim.GetPath()), self._device)] = ( + MultiMeshRayCaster.meshes[base_key] + ) + prim_key = (str(target_prim.GetPath()), self._device) + if prim_key in MultiMeshRayCaster.meshes: + wp_mesh_ids.append(MultiMeshRayCaster.meshes[prim_key].id) loaded_vertices.append(None) continue @@ -219,7 +209,6 @@ def _initialize_warp_meshes(self): trimesh_meshes = [] for mesh_prim in mesh_prims: - # check if valid if mesh_prim is None or not mesh_prim.IsValid(): raise RuntimeError(f"Invalid mesh prim path: {target_prim}") @@ -240,13 +229,11 @@ def _initialize_warp_meshes(self): transform[:3, 3] = relative_pos.numpy() mesh.apply_transform(transform) - # add to list of parsed meshes trimesh_meshes.append(mesh) if len(trimesh_meshes) == 1: trimesh_mesh = trimesh_meshes[0] elif target_cfg.merge_prim_meshes: - # combine all trimesh meshes into a single mesh trimesh_mesh = trimesh.util.concatenate(trimesh_meshes) else: raise RuntimeError( @@ -254,20 +241,17 @@ def _initialize_warp_meshes(self): " enable `merge_prim_meshes` in the configuration or specify each mesh separately." ) - # check if the mesh is already registered, if so only reference the mesh registered_idx = _registered_points_idx(trimesh_mesh.vertices, loaded_vertices) if registered_idx != -1 and self.cfg.reference_meshes: logger.info("Found a duplicate mesh, only reference the mesh.") - # Found a duplicate mesh, only reference the mesh. loaded_vertices.append(None) wp_mesh_ids.append(wp_mesh_ids[registered_idx]) else: loaded_vertices.append(trimesh_mesh.vertices) - wp_mesh = convert_to_warp_mesh(trimesh_mesh.vertices, trimesh_mesh.faces, device=self.device) - MultiMeshRayCaster.meshes[str(target_prim.GetPath())] = wp_mesh + wp_mesh = convert_to_warp_mesh(trimesh_mesh.vertices, trimesh_mesh.faces, device=self._device) + MultiMeshRayCaster.meshes[(str(target_prim.GetPath()), self._device)] = wp_mesh wp_mesh_ids.append(wp_mesh.id) - # print info if registered_idx != -1: logger.info(f"Found duplicate mesh for mesh prims under path '{target_prim.GetPath()}'.") else: @@ -277,12 +261,9 @@ def _initialize_warp_meshes(self): ) if is_global_prim: - # reference the mesh for each environment to ray cast against multi_mesh_ids[target_prim_path] = [wp_mesh_ids] * self._num_envs self._num_meshes_per_env[target_prim_path] = len(wp_mesh_ids) else: - # split up the meshes for each environment. Little bit ugly, since - # the current order is interleaved (env1_obj1, env1_obj2, env2_obj1, env2_obj2, ...) multi_mesh_ids[target_prim_path] = [] mesh_idx = 0 n_meshes_per_env = len(wp_mesh_ids) // self._num_envs @@ -292,26 +273,28 @@ 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 ) - # throw an error if no meshes are found if all([target_cfg.prim_expr not in multi_mesh_ids for target_cfg in self._raycast_targets_cfg]): raise RuntimeError( f"No meshes found for ray-casting! Please check the mesh prim paths: {self.cfg.mesh_prim_paths}" ) total_n_meshes_per_env = sum(self._num_meshes_per_env.values()) - self._mesh_positions_w = torch.zeros(self._num_envs, total_n_meshes_per_env, 3, device=self.device) - self._mesh_orientations_w = torch.zeros(self._num_envs, total_n_meshes_per_env, 4, device=self.device) + self._mesh_positions_w = wp.zeros((self._num_envs, total_n_meshes_per_env), dtype=wp.vec3, device=self.device) + self._mesh_orientations_w = wp.zeros( + (self._num_envs, total_n_meshes_per_env), dtype=wp.quat, device=self.device + ) + # Zero-copy torch views for writing from physics view results (torch tensors) + self._mesh_positions_w_torch = wp.to_torch(self._mesh_positions_w) + self._mesh_orientations_w_torch = wp.to_torch(self._mesh_orientations_w) - # Update the mesh positions and rotations mesh_idx = 0 for target_cfg in self._raycast_targets_cfg: n_meshes = self._num_meshes_per_env[target_cfg.prim_expr] - # update position of the target meshes pos_w, ori_w = [], [] for prim in sim_utils.find_matching_prims(target_cfg.prim_expr): translation, quat = sim_utils.resolve_prim_pose(prim) @@ -320,11 +303,10 @@ def _initialize_warp_meshes(self): pos_w = torch.tensor(pos_w, device=self.device, dtype=torch.float32).view(-1, n_meshes, 3) ori_w = torch.tensor(ori_w, device=self.device, dtype=torch.float32).view(-1, n_meshes, 4) - self._mesh_positions_w[:, mesh_idx : mesh_idx + n_meshes] = pos_w - self._mesh_orientations_w[:, mesh_idx : mesh_idx + n_meshes] = ori_w + self._mesh_positions_w_torch[:, mesh_idx : mesh_idx + n_meshes] = pos_w + self._mesh_orientations_w_torch[:, mesh_idx : mesh_idx + n_meshes] = ori_w mesh_idx += n_meshes - # flatten the list of meshes that are included in mesh_prim_paths of the specific ray caster multi_mesh_ids_flattened = [] for env_idx in range(self._num_envs): meshes_in_env = [] @@ -337,25 +319,30 @@ def _initialize_warp_meshes(self): for target_cfg in self._raycast_targets_cfg ] - # save a warp array with mesh ids that is passed to the raycast function self._mesh_ids_wp = wp.array2d(multi_mesh_ids_flattened, dtype=wp.uint64, device=self.device) def _initialize_rays_impl(self): super()._initialize_rays_impl() + # Persistent buffer for tracking closest-hit distance across meshes (for atomic_min) + self._ray_distance_w = wp.zeros((self._view.count, self.num_rays), dtype=wp.float32, device=self._device) if self.cfg.update_mesh_ids: - self._data.ray_mesh_ids = torch.zeros( - self._num_envs, self.num_rays, 1, device=self.device, dtype=torch.int16 - ) - - def _update_buffers_impl(self, env_mask: wp.array): - """Fills the buffers of the sensor data.""" - env_ids = wp.to_torch(env_mask).nonzero(as_tuple=False).squeeze(-1) - if len(env_ids) == 0: - return - - self._update_ray_infos(env_ids) - - # Update the mesh positions and rotations + self._ray_mesh_id_w = wp.zeros((self._view.count, self.num_rays), dtype=wp.int16, device=self._device) + # Zero-copy torch view with the trailing dim expected by consumers of ray_mesh_ids + self._data.ray_mesh_ids = wp.to_torch(self._ray_mesh_id_w).unsqueeze(-1) + else: + # Dummy 1×1 buffer so the kernel launch always has a valid array to bind + self._ray_mesh_id_w = wp.empty((1, 1), dtype=wp.int16, device=self._device) + # Persistent dummy buffers for unused kernel outputs; allocated once to avoid per-step allocations. + self._dummy_normal_w = wp.empty((1, 1), dtype=wp.vec3, device=self._device) + self._dummy_face_id_w = wp.empty((1, 1), dtype=wp.int32, device=self._device) + + def _update_mesh_transforms(self) -> None: + """Update world-frame mesh positions and orientations for dynamically tracked targets. + + Iterates over all tracked views and writes the current world poses into + ``_mesh_positions_w_torch`` and ``_mesh_orientations_w_torch``. Static (non-tracked) + targets are skipped; their initial poses were set during :meth:`_initialize_warp_meshes`. + """ mesh_idx = 0 for view, target_cfg in zip(self._mesh_views, self._raycast_targets_cfg): if not target_cfg.track_mesh_transforms: @@ -363,42 +350,75 @@ def _update_buffers_impl(self, env_mask: wp.array): continue # update position of the target meshes - pos_w, ori_w = obtain_world_pose_from_view(view, None) + 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 -= pos_offset - ori_w = quat_mul(ori_offset.expand(ori_w.shape[0], -1), ori_w) - count = view.count - if count != 1: # Mesh is not global, i.e. we have different meshes for each env + if count != 1: count = count // self._num_envs pos_w = pos_w.view(self._num_envs, count, 3) ori_w = ori_w.view(self._num_envs, count, 4) - self._mesh_positions_w[:, mesh_idx : mesh_idx + count] = pos_w - self._mesh_orientations_w[:, mesh_idx : mesh_idx + count] = ori_w + self._mesh_positions_w_torch[:, mesh_idx : mesh_idx + count] = pos_w + self._mesh_orientations_w_torch[:, mesh_idx : mesh_idx + count] = ori_w mesh_idx += count - self._data.ray_hits_w[env_ids], _, _, _, mesh_ids = raycast_dynamic_meshes( - self._ray_starts_w[env_ids], - self._ray_directions_w[env_ids], - mesh_ids_wp=self._mesh_ids_wp, # list with shape num_envs x num_meshes_per_env - max_dist=self.cfg.max_distance, - mesh_positions_w=self._mesh_positions_w[env_ids], - mesh_orientations_w=self._mesh_orientations_w[env_ids], - return_mesh_id=self.cfg.update_mesh_ids, + def _update_buffers_impl(self, env_mask: wp.array): + """Fills the buffers of the sensor data.""" + self._update_ray_infos(env_mask) + self._update_mesh_transforms() + + n_meshes = self._mesh_ids_wp.shape[1] + + # Fill output and distance buffers with inf for masked environments + wp.launch( + fill_vec3_inf_kernel, + dim=(self._num_envs, self.num_rays), + inputs=[env_mask, float("inf"), self._data._ray_hits_w], + device=self._device, + ) + wp.launch( + fill_float2d_masked_kernel, + dim=(self._num_envs, self.num_rays), + inputs=[env_mask, float("inf"), self._ray_distance_w], + device=self._device, ) - if self.cfg.update_mesh_ids: - self._data.ray_mesh_ids[env_ids] = mesh_ids + # Ray-cast against all meshes; closest hit wins via atomic_min on ray_distance + wp.launch( + warp_kernels.raycast_dynamic_meshes_kernel, + dim=(n_meshes, self._num_envs, self.num_rays), + inputs=[ + env_mask, + self._mesh_ids_wp, + self._ray_starts_w, + self._ray_directions_w, + self._data._ray_hits_w, + self._ray_distance_w, + self._dummy_normal_w, + self._dummy_face_id_w, + self._ray_mesh_id_w, + self._mesh_positions_w, + self._mesh_orientations_w, + float(self.cfg.max_distance), + int(False), + int(False), + int(self.cfg.update_mesh_ids), + ], + device=self._device, + ) + + def _invalidate_initialize_callback(self, event): + """Invalidates the scene elements.""" + super()._invalidate_initialize_callback(event) + # clear mesh views so they are re-created on the next initialization + MultiMeshRayCaster.mesh_views.clear() 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 a1be3160d99b..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 @@ -5,18 +5,22 @@ from __future__ import annotations -from collections.abc import Sequence from typing import TYPE_CHECKING import torch import warp as wp import isaaclab.utils.math as math_utils -from isaaclab.utils.warp import raycast_dynamic_meshes - +from isaaclab.utils.warp import kernels as warp_kernels + +from .kernels import ( + CAMERA_RAYCAST_MAX_DIST, + compute_distance_to_image_plane_masked_kernel, + fill_float2d_masked_kernel, + fill_vec3_inf_kernel, +) 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: @@ -85,136 +89,205 @@ def _create_buffers(self): ) def _initialize_rays_impl(self): - # Create all indices buffer + # NOTE: This method intentionally does NOT call super()._initialize_rays_impl() through the MRO + # chain. The intermediate classes (RayCasterCamera, MultiMeshRayCaster) use different internal + # buffer names and orderings that are incompatible with the camera's full init path: + # - RayCasterCamera creates single-mesh ray buffers (_ray_distance, _ray_normal_w, etc.) + # - MultiMeshRayCaster creates _ray_distance_w / _ray_mesh_id_w for multi-mesh use + # The camera replaces all of these with its own camera-named equivalents below. + # If either parent class gains new shared buffers, they must be added here explicitly. + + # Camera-specific bookkeeping buffers self._ALL_INDICES = torch.arange(self._view.count, device=self._device, dtype=torch.long) - # Create frame count buffer self._frame = torch.zeros(self._view.count, device=self._device, dtype=torch.long) - # create buffers + + # Build camera output buffers (intrinsics, image data, etc.) self._create_buffers() - # compute intrinsic matrices self._compute_intrinsic_matrices() - # compute ray stars and directions - self.ray_starts, self.ray_directions = self.cfg.pattern_cfg.func( + + # Compute local ray starts/directions from the camera pattern (torch, init-time only) + ray_starts_local, ray_directions_local = self.cfg.pattern_cfg.func( self.cfg.pattern_cfg, self._data.intrinsic_matrices, self._device ) - self.num_rays = self.ray_directions.shape[1] - # create buffer to store ray hits - self.ray_hits_w = torch.zeros(self._view.count, self.num_rays, 3, device=self._device) - # set offsets - quat_w = math_utils.convert_camera_frame_orientation_convention( - torch.tensor([self.cfg.offset.rot], device=self._device), origin=self.cfg.offset.convention, target="world" + self.num_rays = ray_directions_local.shape[1] + + # Store local (sensor-frame) ray arrays as torch tensors for per-env camera-convention rotation + self.ray_starts = ray_starts_local + self.ray_directions = ray_directions_local + + # Camera-frame offset: convert from cfg convention to world convention + quat_offset = math_utils.convert_camera_frame_orientation_convention( + torch.tensor([self.cfg.offset.rot], device=self._device), + origin=self.cfg.offset.convention, + target="world", ) - self._offset_quat = quat_w.repeat(self._view.count, 1) + self._offset_quat = quat_offset.repeat(self._view.count, 1) self._offset_pos = torch.tensor(list(self.cfg.offset.pos), device=self._device).repeat(self._view.count, 1) - self._data.quat_w = torch.zeros(self._view.count, 4, device=self.device) - self._data.pos_w = torch.zeros(self._view.count, 3, device=self.device) + # Camera pose buffers (torch, part of CameraData) + self._data.pos_w = torch.zeros(self._view.count, 3, device=self._device) + self._data.quat_w_world = torch.zeros(self._view.count, 4, device=self._device) + # Warp-backed camera orientation buffer for warp kernel calls; + # updated from self._data.quat_w_world in _update_ray_infos. + self._quat_w_wp = wp.zeros(self._view.count, dtype=wp.quatf, device=self._device) + self._quat_w_wp_torch = wp.to_torch(self._quat_w_wp) + + # Warp buffer for distance_to_image_plane output (if requested) + if "distance_to_image_plane" in self.cfg.data_types: + self._distance_to_image_plane_wp = wp.zeros( + (self._view.count, self.num_rays), dtype=wp.float32, device=self._device + ) + + # World-frame ray buffers: allocate as warp arrays first, then create zero-copy torch views. + # Keeping warp arrays as primary storage avoids lifetime issues when passing to kernels. + self._ray_starts_w = wp.zeros((self._view.count, self.num_rays), dtype=wp.vec3f, device=self._device) + self._ray_directions_w = wp.zeros((self._view.count, self.num_rays), dtype=wp.vec3f, device=self._device) + # Zero-copy torch views used for indexing and post-processing + self._ray_starts_w_torch = wp.to_torch(self._ray_starts_w) + self._ray_directions_w_torch = wp.to_torch(self._ray_directions_w) + + # Ray hit positions as a warp array; expose a torch view for debug visualisation + self._ray_hits_w_cam = wp.zeros((self._view.count, self.num_rays), dtype=wp.vec3f, device=self._device) + self.ray_hits_w = wp.to_torch(self._ray_hits_w_cam) + + # Per-ray closest-hit distance for atomic_min across meshes + self._ray_distance_cam_w = wp.zeros((self._view.count, self.num_rays), dtype=wp.float32, device=self._device) + + # Optional normal buffer (always allocated; filled only when "normals" is requested) + self._ray_normal_w = wp.zeros((self._view.count, self.num_rays), dtype=wp.vec3f, device=self._device) - self._ray_starts_w = torch.zeros(self._view.count, self.num_rays, 3, device=self.device) - self._ray_directions_w = torch.zeros(self._view.count, self.num_rays, 3, device=self.device) + # Mesh-id buffers from MultiMeshRayCaster._initialize_rays_impl + if self.cfg.update_mesh_ids: + self._ray_mesh_id_w = wp.zeros((self._view.count, self.num_rays), dtype=wp.int16, device=self._device) + self._data.ray_mesh_ids = wp.to_torch(self._ray_mesh_id_w).unsqueeze(-1) + else: + self._ray_mesh_id_w = wp.empty((1, 1), dtype=wp.int16, device=self._device) + + # Dummy face-id buffer (not used by camera but required by kernel signature) + self._ray_face_id_w = wp.empty((1, 1), dtype=wp.int32, device=self._device) - def _update_ray_infos(self, env_ids: Sequence[int]): - """Updates the ray information buffers.""" + def _update_ray_infos(self, env_mask: wp.array): + """Updates camera poses and world-frame ray buffers for masked environments. + + Args: + env_mask: Boolean mask selecting which environments to update. Shape is (num_envs,). + """ + env_ids = wp.to_torch(env_mask).nonzero(as_tuple=False).squeeze(-1) + if len(env_ids) == 0: + return - # compute poses from current view - pos_w, quat_w = obtain_world_pose_from_view(self._view, env_ids) + # Compute camera world poses by composing view pose with sensor offset + 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] ) - # update the data + # Store camera pose in CameraData (torch tensors) and warp-backed orientation buffer self._data.pos_w[env_ids] = pos_w self._data.quat_w_world[env_ids] = quat_w - self._data.quat_w_ros[env_ids] = quat_w + self._quat_w_wp_torch[env_ids] = quat_w - # note: full orientation is considered - ray_starts_w = math_utils.quat_apply(quat_w.repeat(1, self.num_rays), self.ray_starts[env_ids]) - ray_starts_w += pos_w.unsqueeze(1) - ray_directions_w = math_utils.quat_apply(quat_w.repeat(1, self.num_rays), self.ray_directions[env_ids]) + # Rotate local ray starts and directions into world frame using full camera orientation + quat_w_repeated = quat_w.repeat(1, self.num_rays).reshape(-1, 4) + ray_starts_local = self.ray_starts[env_ids].reshape(-1, 3) + ray_dirs_local = self.ray_directions[env_ids].reshape(-1, 3) - self._ray_starts_w[env_ids] = ray_starts_w - self._ray_directions_w[env_ids] = ray_directions_w + ray_starts_world = math_utils.quat_apply(quat_w_repeated, ray_starts_local).reshape( + len(env_ids), self.num_rays, 3 + ) + ray_starts_world += pos_w.unsqueeze(1) + ray_dirs_world = math_utils.quat_apply(quat_w_repeated, ray_dirs_local).reshape(len(env_ids), self.num_rays, 3) + + # Write back into the warp-backed buffers via zero-copy torch views + self._ray_starts_w_torch[env_ids] = ray_starts_world + self._ray_directions_w_torch[env_ids] = ray_dirs_world def _update_buffers_impl(self, env_mask: wp.array): """Fills the buffers of the sensor data.""" env_ids = wp.to_torch(env_mask).nonzero(as_tuple=False).squeeze(-1) if len(env_ids) == 0: return - self._update_ray_infos(env_ids) - # increment frame count + self._update_ray_infos(env_mask) + + # Increment frame count for updated environments self._frame[env_ids] += 1 - # Update the mesh positions and rotations - mesh_idx = 0 - for view, target_cfg in zip(self._mesh_views, self._raycast_targets_cfg): - if not target_cfg.track_mesh_transforms: - mesh_idx += self._num_meshes_per_env[target_cfg.prim_expr] - continue - - # update position of the target meshes - pos_w, ori_w = obtain_world_pose_from_view(view, None) - 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 -= pos_offset - ori_w = math_utils.quat_mul(ori_offset.expand(ori_w.shape[0], -1), ori_w) - - count = view.count - if count != 1: # Mesh is not global, i.e. we have different meshes for each env - count = count // self._num_envs - pos_w = pos_w.view(self._num_envs, count, 3) - ori_w = ori_w.view(self._num_envs, count, 4) - - self._mesh_positions_w[:, mesh_idx : mesh_idx + count] = pos_w - self._mesh_orientations_w[:, mesh_idx : mesh_idx + count] = ori_w - mesh_idx += count - - # ray cast and store the hits - self.ray_hits_w[env_ids], ray_depth, ray_normal, _, ray_mesh_ids = raycast_dynamic_meshes( - self._ray_starts_w[env_ids], - self._ray_directions_w[env_ids], - mesh_ids_wp=self._mesh_ids_wp, # list with shape num_envs x num_meshes_per_env - max_dist=self.cfg.max_distance, - mesh_positions_w=self._mesh_positions_w[env_ids], - mesh_orientations_w=self._mesh_orientations_w[env_ids], - return_distance=any( - [name in self.cfg.data_types for name in ["distance_to_image_plane", "distance_to_camera"]] - ), - return_normal="normals" in self.cfg.data_types, - return_mesh_id=self.cfg.update_mesh_ids, + self._update_mesh_transforms() + + n_meshes = self._mesh_ids_wp.shape[1] + return_normal = "normals" in self.cfg.data_types + + # Fill ray hit and distance buffers with inf for masked environments + wp.launch( + fill_vec3_inf_kernel, + dim=(self._num_envs, self.num_rays), + inputs=[env_mask, float("inf"), self._ray_hits_w_cam], + device=self._device, + ) + wp.launch( + fill_float2d_masked_kernel, + dim=(self._num_envs, self.num_rays), + inputs=[env_mask, float("inf"), self._ray_distance_cam_w], + device=self._device, + ) + if return_normal: + wp.launch( + fill_vec3_inf_kernel, + dim=(self._num_envs, self.num_rays), + inputs=[env_mask, float("inf"), self._ray_normal_w], + device=self._device, + ) + + # Ray-cast against all meshes; closest hit wins via atomic_min on ray_distance + wp.launch( + warp_kernels.raycast_dynamic_meshes_kernel, + dim=(n_meshes, self._num_envs, self.num_rays), + inputs=[ + env_mask, + self._mesh_ids_wp, + self._ray_starts_w, + self._ray_directions_w, + self._ray_hits_w_cam, + self._ray_distance_cam_w, + self._ray_normal_w, + self._ray_face_id_w, + self._ray_mesh_id_w, + self._mesh_positions_w, + self._mesh_orientations_w, + float(CAMERA_RAYCAST_MAX_DIST), + int(return_normal), + int(False), + int(self.cfg.update_mesh_ids), + ], + device=self._device, ) - # update output buffers if "distance_to_image_plane" in self.cfg.data_types: - # note: data is in camera frame so we only take the first component (z-axis of camera frame) - distance_to_image_plane = ( - math_utils.quat_apply( - math_utils.quat_inv(self._data.quat_w_world[env_ids]).repeat(1, self.num_rays), - (ray_depth[:, :, None] * self._ray_directions_w[env_ids]), - ) - )[:, :, 0] - # apply the maximum distance after the transformation - if self.cfg.depth_clipping_behavior == "max": - distance_to_image_plane = torch.clip(distance_to_image_plane, max=self.cfg.max_distance) - distance_to_image_plane[torch.isnan(distance_to_image_plane)] = self.cfg.max_distance - elif self.cfg.depth_clipping_behavior == "zero": - distance_to_image_plane[distance_to_image_plane > self.cfg.max_distance] = 0.0 - distance_to_image_plane[torch.isnan(distance_to_image_plane)] = 0.0 - self._data.output["distance_to_image_plane"][env_ids] = distance_to_image_plane.view( - -1, *self.image_shape, 1 + wp.launch( + compute_distance_to_image_plane_masked_kernel, + dim=(self._num_envs, self.num_rays), + inputs=[env_mask, self._quat_w_wp, self._ray_distance_cam_w, self._ray_directions_w], + outputs=[self._distance_to_image_plane_wp], + device=self._device, ) + # Apply depth clipping on the intermediate buffer (leaves _ray_distance_cam_w unmodified) + self._apply_depth_clipping(env_mask, self._distance_to_image_plane_wp) + d2ip_torch = wp.to_torch(self._distance_to_image_plane_wp) + self._data.output["distance_to_image_plane"][env_ids] = d2ip_torch[env_ids].view(-1, *self.image_shape, 1) if "distance_to_camera" in self.cfg.data_types: - if self.cfg.depth_clipping_behavior == "max": - ray_depth = torch.clip(ray_depth, max=self.cfg.max_distance) - elif self.cfg.depth_clipping_behavior == "zero": - ray_depth[ray_depth > self.cfg.max_distance] = 0.0 - self._data.output["distance_to_camera"][env_ids] = ray_depth.view(-1, *self.image_shape, 1) + # d2ip (if requested) was computed before this block so _ray_distance_cam_w is still unclipped. + self._apply_depth_clipping(env_mask, self._ray_distance_cam_w) + ray_dist_torch = wp.to_torch(self._ray_distance_cam_w) + self._data.output["distance_to_camera"][env_ids] = ray_dist_torch[env_ids].view(-1, *self.image_shape, 1) - if "normals" in self.cfg.data_types: - self._data.output["normals"][env_ids] = ray_normal.view(-1, *self.image_shape, 3) + if return_normal: + ray_normal_torch = wp.to_torch(self._ray_normal_w) + self._data.output["normals"][env_ids] = ray_normal_torch[env_ids].view(-1, *self.image_shape, 3) if self.cfg.update_mesh_ids: - self._data.image_mesh_ids[env_ids] = ray_mesh_ids.view(-1, *self.image_shape, 1) + self._data.image_mesh_ids[env_ids] = wp.to_torch(self._ray_mesh_id_w)[env_ids].view( + -1, *self.image_shape, 1 + ) diff --git a/source/isaaclab/isaaclab/sensors/ray_caster/multi_mesh_ray_caster_camera_data.py b/source/isaaclab/isaaclab/sensors/ray_caster/multi_mesh_ray_caster_camera_data.py index d2f26abdbf47..21338f0a0616 100644 --- a/source/isaaclab/isaaclab/sensors/ray_caster/multi_mesh_ray_caster_camera_data.py +++ b/source/isaaclab/isaaclab/sensors/ray_caster/multi_mesh_ray_caster_camera_data.py @@ -9,11 +9,15 @@ from isaaclab.sensors.camera import CameraData -from .ray_caster_data import RayCasterData +class MultiMeshRayCasterCameraData(CameraData): + """Data container for the multi-mesh ray-cast camera sensor. -class MultiMeshRayCasterCameraData(CameraData, RayCasterData): - """Data container for the multi-mesh ray-cast sensor.""" + This class extends :class:`CameraData` with additional mesh-id information. + It does not inherit from :class:`RayCasterData` because the camera variant + manages its own torch-based pose and hit buffers independently from the + warp-native :class:`RayCasterData`. + """ image_mesh_ids: torch.Tensor = None """The mesh ids of the image pixels. 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 731d57f1638f..1e23ee00c1b6 100644 --- a/source/isaaclab/isaaclab/sensors/ray_caster/ray_caster.py +++ b/source/isaaclab/isaaclab/sensors/ray_caster/ray_caster.py @@ -13,24 +13,27 @@ import torch import warp as wp -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.math import quat_apply, quat_apply_yaw -from isaaclab.utils.warp import convert_to_warp_mesh, raycast_mesh +from isaaclab.utils.warp import convert_to_warp_mesh +from isaaclab.utils.warp.kernels import raycast_mesh_masked_kernel from ..sensor_base import SensorBase -from .ray_cast_utils import obtain_world_pose_from_view +from .kernels import ( + apply_z_drift_kernel, + fill_vec3_inf_kernel, + update_ray_caster_kernel, +) from .ray_caster_data import RayCasterData if TYPE_CHECKING: from .ray_caster_cfg import RayCasterCfg -# import logger logger = logging.getLogger(__name__) @@ -42,8 +45,8 @@ class RayCaster(SensorBase): a set of meshes with a given ray pattern. The meshes are parsed from the list of primitive paths provided in the configuration. These are then - converted to warp meshes and stored in the `warp_meshes` list. The ray-caster then ray-casts against - these warp meshes using the ray pattern provided in the configuration. + converted to warp meshes and stored in the :attr:`meshes` dictionary. The ray-caster then ray-casts + against these warp meshes using the ray pattern provided in the configuration. .. note:: Currently, only static meshes are supported. Extending the warp mesh to support dynamic meshes @@ -53,11 +56,12 @@ class RayCaster(SensorBase): cfg: RayCasterCfg """The configuration parameters.""" - # Class variables to share meshes across instances - meshes: ClassVar[dict[str, wp.Mesh]] = {} + meshes: ClassVar[dict[tuple[str, str], wp.Mesh]] = {} """A dictionary to store warp meshes for raycasting, shared across all instances. - The keys correspond to the prim path for the meshes, and values are the corresponding warp Mesh objects.""" + The keys are ``(prim_path, device)`` tuples and values are the corresponding warp Mesh objects. + Including the device in the key prevents a mesh created on one device (e.g. CPU) from being + reused by a kernel running on a different device (e.g. CUDA).""" _instance_count: ClassVar[int] = 0 """A counter to track the number of RayCaster instances, used to manage class variable lifecycle.""" @@ -68,9 +72,9 @@ def __init__(self, cfg: RayCasterCfg): cfg: The configuration parameters. """ RayCaster._instance_count += 1 - # Initialize base class super().__init__(cfg) - # Create empty variables for storing output data + # Resolve physics-body paths and spawn the sensor Xform child if needed. + self._resolve_and_spawn("raycaster") self._data = RayCasterData() def __str__(self) -> str: @@ -116,10 +120,10 @@ def reset(self, env_ids: Sequence[int] | None = None, env_mask: wp.array | None else: env_ids = slice(None) num_envs_ids = self._view.count - # resample the drift + # resample drift (uses torch views for indexing) r = torch.empty(num_envs_ids, 3, device=self.device) self.drift[env_ids] = r.uniform_(*self.cfg.drift_range) - # resample the height drift + # resample the ray cast drift range_list = [self.cfg.ray_cast_drift_range.get(key, (0.0, 0.0)) for key in ["x", "y", "z"]] ranges = torch.tensor(range_list, device=self.device) self.ray_cast_drift[env_ids] = math_utils.sample_uniform( @@ -132,21 +136,28 @@ 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) + # 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) + + # 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: + raise RuntimeError(f"Unsupported ray_alignment type: {self.cfg.ray_alignment}.") + self._alignment_mode = alignment_map[self.cfg.ray_alignment] # load the meshes by parsing the stage self._initialize_warp_meshes() - # initialize the ray start and directions self._initialize_rays_impl() def _initialize_warp_meshes(self): @@ -158,168 +169,188 @@ def _initialize_warp_meshes(self): # read prims to ray-cast for mesh_prim_path in self.cfg.mesh_prim_paths: - # check if mesh already casted into warp mesh - if mesh_prim_path in RayCaster.meshes: + mesh_key = (mesh_prim_path, self._device) + if mesh_key in RayCaster.meshes: continue - # check if the prim is a plane - handle PhysX plane as a special case - # if a plane exists then we need to create an infinite mesh that is a plane mesh_prim = sim_utils.get_first_matching_child_prim( mesh_prim_path, lambda prim: prim.GetTypeName() == "Plane" ) - # if we did not find a plane then we need to read the mesh if mesh_prim is None: - # obtain the mesh prim mesh_prim = sim_utils.get_first_matching_child_prim( mesh_prim_path, lambda prim: prim.GetTypeName() == "Mesh" ) - # check if valid if mesh_prim is None or not mesh_prim.IsValid(): raise RuntimeError(f"Invalid mesh prim path: {mesh_prim_path}") - # cast into UsdGeomMesh mesh_prim = UsdGeom.Mesh(mesh_prim) - # read the vertices and faces points = np.asarray(mesh_prim.GetPointsAttr().Get()) - # Get world transform using pure USD (UsdGeom.Xformable) xformable = UsdGeom.Xformable(mesh_prim) world_transform: Gf.Matrix4d = xformable.ComputeLocalToWorldTransform(Usd.TimeCode.Default()) transform_matrix = np.array(world_transform).T points = np.matmul(points, transform_matrix[:3, :3].T) points += transform_matrix[:3, 3] indices = np.asarray(mesh_prim.GetFaceVertexIndicesAttr().Get()) - wp_mesh = convert_to_warp_mesh(points, indices, device=self.device) - # print info + wp_mesh = convert_to_warp_mesh(points, indices, device=self._device) logger.info( f"Read mesh prim: {mesh_prim.GetPath()} with {len(points)} vertices and {len(indices)} faces." ) else: mesh = make_plane(size=(2e6, 2e6), height=0.0, center_zero=True) - wp_mesh = convert_to_warp_mesh(mesh.vertices, mesh.faces, device=self.device) - # print info + wp_mesh = convert_to_warp_mesh(mesh.vertices, mesh.faces, device=self._device) logger.info(f"Created infinite plane mesh prim: {mesh_prim.GetPath()}.") - # add the warp mesh to the list - RayCaster.meshes[mesh_prim_path] = wp_mesh + RayCaster.meshes[mesh_key] = wp_mesh - # throw an error if no meshes are found - if all([mesh_prim_path not in RayCaster.meshes for mesh_prim_path in self.cfg.mesh_prim_paths]): + if all((mesh_prim_path, self._device) not in RayCaster.meshes for mesh_prim_path in self.cfg.mesh_prim_paths): raise RuntimeError( f"No meshes found for ray-casting! Please check the mesh prim paths: {self.cfg.mesh_prim_paths}" ) def _initialize_rays_impl(self): - # compute ray stars and directions - self.ray_starts, self.ray_directions = self.cfg.pattern_cfg.func(self.cfg.pattern_cfg, self._device) - self.num_rays = len(self.ray_directions) - # apply offset transformation to the rays + # Compute ray starts and directions from pattern (torch, init-time only) + ray_starts_torch, ray_directions_torch = self.cfg.pattern_cfg.func(self.cfg.pattern_cfg, self._device) + self.num_rays = len(ray_directions_torch) + + # Apply sensor offset rotation/position to local ray pattern offset_pos = torch.tensor(list(self.cfg.offset.pos), device=self._device) offset_quat = torch.tensor(list(self.cfg.offset.rot), device=self._device) - self.ray_directions = quat_apply(offset_quat.repeat(len(self.ray_directions), 1), self.ray_directions) - self.ray_starts += offset_pos - # repeat the rays for each sensor - self.ray_starts = self.ray_starts.repeat(self._view.count, 1, 1) - self.ray_directions = self.ray_directions.repeat(self._view.count, 1, 1) - # prepare drift - self.drift = torch.zeros(self._view.count, 3, device=self.device) - self.ray_cast_drift = torch.zeros(self._view.count, 3, device=self.device) - # fill the data buffer - self._data.pos_w = torch.zeros(self._view.count, 3, device=self.device) - self._data.quat_w = torch.zeros(self._view.count, 4, device=self.device) - self._data.ray_hits_w = torch.zeros(self._view.count, self.num_rays, 3, device=self.device) - self._ray_starts_w = torch.zeros(self._view.count, self.num_rays, 3, device=self.device) - self._ray_directions_w = torch.zeros(self._view.count, self.num_rays, 3, device=self.device) - - def _update_ray_infos(self, env_ids: Sequence[int]): - """Updates the ray information buffers.""" - - pos_w, quat_w = obtain_world_pose_from_view(self._view, env_ids) - pos_w, quat_w = math_utils.combine_frame_transforms( - pos_w, quat_w, self._offset[0][env_ids], self._offset[1][env_ids] + ray_directions_torch = math_utils.quat_apply( + offset_quat.repeat(len(ray_directions_torch), 1), ray_directions_torch ) - # apply drift to ray starting position in world frame - pos_w += self.drift[env_ids] - # store the poses - self._data.pos_w[env_ids] = pos_w - self._data.quat_w[env_ids] = quat_w - - # check if user provided attach_yaw_only flag - 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." - ) - # set ray alignment to yaw - 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'." - # log the warning - logger.warning(msg) - # ray cast based on the sensor poses - if self.cfg.ray_alignment == "world": - # apply horizontal drift to ray starting position in ray caster frame - pos_w[:, 0:2] += self.ray_cast_drift[env_ids, 0:2] - # no rotation is considered and directions are not rotated - ray_starts_w = self.ray_starts[env_ids] - ray_starts_w += pos_w.unsqueeze(1) - ray_directions_w = self.ray_directions[env_ids] - elif self.cfg.ray_alignment == "yaw": - # apply horizontal drift to ray starting position in ray caster frame - pos_w[:, 0:2] += quat_apply_yaw(quat_w, self.ray_cast_drift[env_ids])[:, 0:2] - # only yaw orientation is considered and directions are not rotated - ray_starts_w = quat_apply_yaw(quat_w.repeat(1, self.num_rays), self.ray_starts[env_ids]) - ray_starts_w += pos_w.unsqueeze(1) - ray_directions_w = self.ray_directions[env_ids] - elif self.cfg.ray_alignment == "base": - # apply horizontal drift to ray starting position in ray caster frame - pos_w[:, 0:2] += quat_apply(quat_w, self.ray_cast_drift[env_ids])[:, 0:2] - # full orientation is considered - ray_starts_w = quat_apply(quat_w.repeat(1, self.num_rays), self.ray_starts[env_ids]) - ray_starts_w += pos_w.unsqueeze(1) - ray_directions_w = quat_apply(quat_w.repeat(1, self.num_rays), self.ray_directions[env_ids]) - else: - raise RuntimeError(f"Unsupported ray_alignment type: {self.cfg.ray_alignment}.") + ray_starts_torch += offset_pos + + # Repeat for each environment + ray_starts_torch = ray_starts_torch.repeat(self._view.count, 1, 1) + ray_directions_torch = ray_directions_torch.repeat(self._view.count, 1, 1) + + # Create warp arrays from the init-time torch data + # The warp arrays own the memory; torch views provide backward-compat indexing + self._ray_starts_local = wp.from_torch(ray_starts_torch.contiguous(), dtype=wp.vec3f) + self._ray_directions_local = wp.from_torch(ray_directions_torch.contiguous(), dtype=wp.vec3f) - self._ray_starts_w[env_ids] = ray_starts_w - self._ray_directions_w[env_ids] = ray_directions_w + # Torch views (same attribute names as before for subclass compatibility) + self.ray_starts = wp.to_torch(self._ray_starts_local) + self.ray_directions = wp.to_torch(self._ray_directions_local) + + # Drift buffers (warp-owned, torch views for reset indexing) + self._drift = wp.zeros(self._view.count, dtype=wp.vec3f, device=self._device) + self._ray_cast_drift = wp.zeros(self._view.count, dtype=wp.vec3f, device=self._device) + self.drift = wp.to_torch(self._drift) + self.ray_cast_drift = wp.to_torch(self._ray_cast_drift) + + # World-frame ray buffers + self._ray_starts_w = wp.zeros((self._view.count, self.num_rays), dtype=wp.vec3f, device=self._device) + self._ray_directions_w = wp.zeros((self._view.count, self.num_rays), dtype=wp.vec3f, device=self._device) + + # Torch views for subclass compatibility + self._ray_starts_w_torch = wp.to_torch(self._ray_starts_w) + self._ray_directions_w_torch = wp.to_torch(self._ray_directions_w) + + # Data buffers + self._data.create_buffers(self._view.count, self.num_rays, self._device) + + # Dummy distance/normal buffers required by the merged raycast_mesh_masked_kernel signature. + # Sized (1, 1) even though the kernel is launched at (num_envs, num_rays): the kernel only + # writes to these buffers when return_distance==1 or return_normal==1 respectively, and + # RayCaster always passes 0 for both flags. If those flags are ever enabled here, these + # buffers must be resized to (num_envs, num_rays) to avoid an out-of-bounds write. + self._dummy_ray_distance = wp.empty((1, 1), dtype=wp.float32, device=self._device) + 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 frame view as a warp array of ``wp.transformf``. + + Returns: + 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`. + """ + 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.""" + transforms = self._get_view_transforms_wp() + + wp.launch( + update_ray_caster_kernel, + dim=(self._num_envs, self.num_rays), + inputs=[ + transforms, + env_mask, + self._offset_pos_wp, + self._offset_quat_wp, + self._drift, + self._ray_cast_drift, + self._ray_starts_local, + self._ray_directions_local, + self._alignment_mode, + ], + outputs=[ + self._data._pos_w, + self._data._quat_w, + self._ray_starts_w, + self._ray_directions_w, + ], + device=self._device, + ) def _update_buffers_impl(self, env_mask: wp.array): """Fills the buffers of the sensor data.""" - env_ids = wp.to_torch(env_mask).nonzero(as_tuple=False).squeeze(-1) - if len(env_ids) == 0: - return - self._update_ray_infos(env_ids) + self._update_ray_infos(env_mask) + + # Fill ray hits with inf before raycasting + wp.launch( + fill_vec3_inf_kernel, + dim=(self._num_envs, self.num_rays), + inputs=[env_mask, float("inf"), self._data._ray_hits_w], + device=self._device, + ) - # ray cast and store the hits - # TODO: Make this work for multiple meshes? - self._data.ray_hits_w[env_ids] = raycast_mesh( - self._ray_starts_w[env_ids], - self._ray_directions_w[env_ids], - max_dist=self.cfg.max_distance, - mesh=RayCaster.meshes[self.cfg.mesh_prim_paths[0]], - )[0] + # Ray-cast against the mesh + wp.launch( + raycast_mesh_masked_kernel, + dim=(self._num_envs, self.num_rays), + inputs=[ + RayCaster.meshes[(self.cfg.mesh_prim_paths[0], self._device)].id, + env_mask, + self._ray_starts_w, + self._ray_directions_w, + float(self.cfg.max_distance), + int(False), # return_distance: not needed by RayCaster + int(False), # return_normal: not needed by RayCaster + self._data._ray_hits_w, + self._dummy_ray_distance, + self._dummy_ray_normal, + ], + device=self._device, + ) - # apply vertical drift to ray starting position in ray caster frame - self._data.ray_hits_w[env_ids, :, 2] += self.ray_cast_drift[env_ids, 2].unsqueeze(-1) + # Apply vertical drift to ray hits + wp.launch( + apply_z_drift_kernel, + dim=(self._num_envs, self.num_rays), + inputs=[env_mask, self._ray_cast_drift, self._data._ray_hits_w], + device=self._device, + ) def _set_debug_vis_impl(self, debug_vis: bool): - # set visibility of markers - # note: parent only deals with callbacks. not their visibility if debug_vis: if not hasattr(self, "ray_visualizer"): self.ray_visualizer = VisualizationMarkers(self.cfg.visualizer_cfg) - # set their visibility to true self.ray_visualizer.set_visibility(True) else: if hasattr(self, "ray_visualizer"): self.ray_visualizer.set_visibility(False) def _debug_vis_callback(self, event): - if self._data.ray_hits_w is None: + if self._data._ray_hits_w is None: return + ray_hits_torch = wp.to_torch(self._data._ray_hits_w) # remove possible inf values - viz_points = self._data.ray_hits_w.reshape(-1, 3) + viz_points = ray_hits_torch.reshape(-1, 3) viz_points = viz_points[~torch.any(torch.isinf(viz_points), dim=1)] # if no points to visualize, skip @@ -328,98 +359,13 @@ 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 | any, tuple[torch.Tensor, torch.Tensor]]: - """Obtain a prim view that can be used to track the pose of the parget 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: - # TODO: Need to handle the case where API is present but it is disabled - 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 - - # TODO: Need to handle the case where API is present but it is disabled - 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 - - # switch the current prim to the parent prim - 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. """ def _invalidate_initialize_callback(self, event): """Invalidates the scene elements.""" - # call parent super()._invalidate_initialize_callback(event) - # set all existing views to None to invalidate them self._view = None def __del__(self): 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 c27f470dcfc0..17bb1e601980 100644 --- a/source/isaaclab/isaaclab/sensors/ray_caster/ray_caster_camera.py +++ b/source/isaaclab/isaaclab/sensors/ray_caster/ray_caster_camera.py @@ -16,9 +16,17 @@ import isaaclab.utils.math as math_utils from isaaclab.sensors.camera import CameraData -from isaaclab.utils.warp import raycast_mesh - -from .ray_cast_utils import obtain_world_pose_from_view +from isaaclab.utils.warp.kernels import raycast_mesh_masked_kernel + +from .kernels import ( + ALIGNMENT_BASE, + CAMERA_RAYCAST_MAX_DIST, + apply_depth_clipping_masked_kernel, + compute_distance_to_image_plane_masked_kernel, + fill_float2d_masked_kernel, + fill_vec3_inf_kernel, + update_ray_caster_kernel, +) from .ray_caster import RayCaster if TYPE_CHECKING: @@ -143,6 +151,13 @@ def set_intrinsic_matrices( self.ray_starts[env_ids], self.ray_directions[env_ids] = self.cfg.pattern_cfg.func( self.cfg.pattern_cfg, self._data.intrinsic_matrices[env_ids], self._device ) + # Refresh warp views of local ray buffers; .contiguous() may produce a copy so we store + # the contiguous tensors explicitly to prevent GC while the warp views are alive. + if hasattr(self, "_ray_starts_local"): + self._ray_starts_contiguous = self.ray_starts.contiguous() + self._ray_directions_contiguous = self.ray_directions.contiguous() + self._ray_starts_local = wp.from_torch(self._ray_starts_contiguous, dtype=wp.vec3f) + self._ray_directions_local = wp.from_torch(self._ray_directions_contiguous, dtype=wp.vec3f) def reset(self, env_ids: Sequence[int] | None = None, env_mask: wp.array | None = None): # reset the timestamps @@ -152,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] ) @@ -179,7 +198,7 @@ def set_world_poses( - :obj:`"ros"` - forward axis: +Z - up axis -Y - Offset is applied in the ROS convention - :obj:`"world"` - forward axis: +X - up axis +Z - Offset is applied in the World Frame convention - See :meth:`isaaclab.utils.maths.convert_camera_frame_orientation_convention` for more details + See :meth:`isaaclab.utils.math.convert_camera_frame_orientation_convention` for more details on the conventions. Args: @@ -198,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 @@ -211,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] ) @@ -224,7 +246,7 @@ def set_world_poses_from_view( """Set the poses of the camera from the eye position and look-at target position. Args: - eyes: The positions of the camera's eye. Shape is N, 3). + eyes: The positions of the camera's eye. Shape is (N, 3). targets: The target locations to look at. Shape is (N, 3). env_ids: A sensor ids to manipulate. Defaults to None, which means all sensor indices. @@ -253,100 +275,243 @@ def _initialize_rays_impl(self): self._create_buffers() # compute intrinsic matrices self._compute_intrinsic_matrices() - # compute ray stars and directions + # compute ray starts and directions self.ray_starts, self.ray_directions = self.cfg.pattern_cfg.func( self.cfg.pattern_cfg, self._data.intrinsic_matrices, self._device ) self.num_rays = self.ray_directions.shape[1] - # create buffer to store ray hits - self.ray_hits_w = torch.zeros(self._view.count, self.num_rays, 3, device=self._device) - # set offsets + + # Offset buffers: warp-primary so the kernel always sees the current values without re-wrapping. + # Zero-copy torch views (_offset_pos, _offset_quat) are used by set_world_poses for indexed writes. + self._offset_pos_wp = wp.zeros(self._view.count, dtype=wp.vec3f, device=self._device) + self._offset_quat_wp = wp.zeros(self._view.count, dtype=wp.quatf, device=self._device) + self._offset_pos = wp.to_torch(self._offset_pos_wp) + self._offset_quat = wp.to_torch(self._offset_quat_wp) + # Initialize from config quat_w = math_utils.convert_camera_frame_orientation_convention( torch.tensor([self.cfg.offset.rot], device=self._device), origin=self.cfg.offset.convention, target="world" ) - self._offset_quat = quat_w.repeat(self._view.count, 1) - self._offset_pos = torch.tensor(list(self.cfg.offset.pos), device=self._device).repeat(self._view.count, 1) + self._offset_pos[:] = torch.tensor(list(self.cfg.offset.pos), device=self._device) + self._offset_quat[:] = quat_w + + # Warp buffers for world-frame rays (used by update kernel) + self._ray_starts_w = wp.zeros((self._view.count, self.num_rays), dtype=wp.vec3f, device=self._device) + self._ray_directions_w = wp.zeros((self._view.count, self.num_rays), dtype=wp.vec3f, device=self._device) + + # Warp views for ray_starts and ray_directions (from torch tensors returned by pattern_cfg.func) + # These are (num_envs, num_rays, 3) torch tensors; wrap as warp vec3f arrays. + # Store contiguous tensors explicitly so they are not garbage-collected while the + # warp views are alive (mirrors the pattern in RayCaster._initialize_impl). + self._ray_starts_contiguous = self.ray_starts.contiguous() + self._ray_directions_contiguous = self.ray_directions.contiguous() + self._ray_starts_local = wp.from_torch(self._ray_starts_contiguous, dtype=wp.vec3f) + self._ray_directions_local = wp.from_torch(self._ray_directions_contiguous, dtype=wp.vec3f) + + # Wrap the torch drift buffers (created in _create_buffers) as warp arrays (zero-copy). + # Cameras do not apply positional drift, so these remain zero. + self._drift_contiguous = self.drift.contiguous() + self._ray_cast_drift_contiguous = self.ray_cast_drift.contiguous() + self._drift = wp.from_torch(self._drift_contiguous, dtype=wp.vec3f) + self._ray_cast_drift = wp.from_torch(self._ray_cast_drift_contiguous, dtype=wp.vec3f) + + # Warp buffers for camera pose outputs + self._pos_w_wp = wp.zeros(self._view.count, dtype=wp.vec3f, device=self._device) + self._quat_w_wp = wp.zeros(self._view.count, dtype=wp.quatf, device=self._device) + + # Intermediate warp buffers for ray results (filled with inf before each raycasting step) + self._ray_distance = wp.zeros((self._view.count, self.num_rays), dtype=wp.float32, device=self._device) + if "normals" in self.cfg.data_types: + self._ray_normal_w = wp.zeros((self._view.count, self.num_rays), dtype=wp.vec3f, device=self._device) + else: + self._ray_normal_w = wp.zeros((1, 1), dtype=wp.vec3f, device=self._device) + + if "distance_to_image_plane" in self.cfg.data_types: + self._distance_to_image_plane_wp = wp.zeros( + (self._view.count, self.num_rays), dtype=wp.float32, device=self._device + ) + + # Torch buffer for ray hits (used by debug visualizer) + self.ray_hits_w = torch.full((self._view.count, self.num_rays, 3), float("inf"), device=self._device) + # Warp view of ray_hits_w + self._ray_hits_w_wp = wp.from_torch(self.ray_hits_w.contiguous(), dtype=wp.vec3f) + + # Cache zero-copy torch views of warp output buffers to avoid per-step wrapper allocation. + self._pos_w_torch = wp.to_torch(self._pos_w_wp) + self._quat_w_torch = wp.to_torch(self._quat_w_wp) + self._ray_distance_torch = wp.to_torch(self._ray_distance) + if "distance_to_image_plane" in self.cfg.data_types: + self._distance_to_image_plane_torch = wp.to_torch(self._distance_to_image_plane_wp) + if "normals" in self.cfg.data_types: + self._ray_normal_w_torch = wp.to_torch(self._ray_normal_w) def _update_buffers_impl(self, env_mask: wp.array): """Fills the buffers of the sensor data.""" + # Convert mask to indices for torch-indexed writes env_ids = wp.to_torch(env_mask).nonzero(as_tuple=False).squeeze(-1) if len(env_ids) == 0: return # increment frame count self._frame[env_ids] += 1 - # compute poses from current view - pos_w, quat_w = obtain_world_pose_from_view(self._view, env_ids, clone=True) - pos_w, quat_w = math_utils.combine_frame_transforms( - pos_w, quat_w, self._offset_pos[env_ids], self._offset_quat[env_ids] + # Update world-frame ray starts/directions and camera pose via warp kernel. + # Camera always uses ALIGNMENT_BASE (full orientation) and zero drift. + transforms = self._get_view_transforms_wp() + wp.launch( + update_ray_caster_kernel, + dim=(self._num_envs, self.num_rays), + inputs=[ + transforms, + env_mask, + self._offset_pos_wp, + self._offset_quat_wp, + self._drift, + self._ray_cast_drift, + self._ray_starts_local, + self._ray_directions_local, + int(ALIGNMENT_BASE), + ], + outputs=[ + self._pos_w_wp, + self._quat_w_wp, + self._ray_starts_w, + self._ray_directions_w, + ], + device=self._device, ) - # update the data - self._data.pos_w[env_ids] = pos_w - self._data.quat_w_world[env_ids] = quat_w - # note: full orientation is considered - ray_starts_w = math_utils.quat_apply(quat_w.repeat(1, self.num_rays), self.ray_starts[env_ids]) - ray_starts_w += pos_w.unsqueeze(1) - ray_directions_w = math_utils.quat_apply(quat_w.repeat(1, self.num_rays), self.ray_directions[env_ids]) - - # ray cast and store the hits - # note: we set max distance to 1e6 during the ray-casting. THis is because we clip the distance - # to the image plane and distance to the camera to the maximum distance afterwards in-order to - # match the USD camera behavior. - - # TODO: Make ray-casting work for multiple meshes? - # necessary for regular dictionaries. - self.ray_hits_w, ray_depth, ray_normal, _ = raycast_mesh( - ray_starts_w, - ray_directions_w, - mesh=RayCaster.meshes[self.cfg.mesh_prim_paths[0]], - max_dist=1e6, - return_distance=any( - [name in self.cfg.data_types for name in ["distance_to_image_plane", "distance_to_camera"]] - ), - return_normal="normals" in self.cfg.data_types, + # Write camera pose to CameraData (torch tensors) + self._data.pos_w[env_ids] = self._pos_w_torch[env_ids] + self._data.quat_w_world[env_ids] = self._quat_w_torch[env_ids] + + # Fill ray hit positions with inf before raycasting + wp.launch( + fill_vec3_inf_kernel, + dim=(self._num_envs, self.num_rays), + inputs=[env_mask, float("inf"), self._ray_hits_w_wp], + device=self._device, + ) + + # Fill ray distance with inf before raycasting + wp.launch( + fill_float2d_masked_kernel, + dim=(self._num_envs, self.num_rays), + inputs=[env_mask, float("inf"), self._ray_distance], + device=self._device, + ) + + # Determine whether to compute normals + need_normal = int("normals" in self.cfg.data_types) + if need_normal: + # Fill normal buffer with inf before raycasting + wp.launch( + fill_vec3_inf_kernel, + dim=(self._num_envs, self.num_rays), + inputs=[env_mask, float("inf"), self._ray_normal_w], + device=self._device, + ) + + # Ray-cast against the mesh; use a large upper-bound max_dist so depth clipping + # can be applied per-data-type afterwards (matching the original behaviour). + wp.launch( + raycast_mesh_masked_kernel, + dim=(self._num_envs, self.num_rays), + inputs=[ + RayCaster.meshes[(self.cfg.mesh_prim_paths[0], self._device)].id, + env_mask, + self._ray_starts_w, + self._ray_directions_w, + float(CAMERA_RAYCAST_MAX_DIST), + int(True), # return_distance: always needed for depth output + need_normal, + self._ray_hits_w_wp, + self._ray_distance, + self._ray_normal_w, + ], + device=self._device, ) - # update output buffers + + # Compute distance_to_image_plane using a warp kernel if "distance_to_image_plane" in self.cfg.data_types: - # note: data is in camera frame so we only take the first component (z-axis of camera frame) - distance_to_image_plane = ( - math_utils.quat_apply( - math_utils.quat_inv(quat_w).repeat(1, self.num_rays), - (ray_depth[:, :, None] * ray_directions_w), - ) - )[:, :, 0] - # apply the maximum distance after the transformation - if self.cfg.depth_clipping_behavior == "max": - distance_to_image_plane = torch.clip(distance_to_image_plane, max=self.cfg.max_distance) - distance_to_image_plane[torch.isnan(distance_to_image_plane)] = self.cfg.max_distance - elif self.cfg.depth_clipping_behavior == "zero": - distance_to_image_plane[distance_to_image_plane > self.cfg.max_distance] = 0.0 - distance_to_image_plane[torch.isnan(distance_to_image_plane)] = 0.0 - self._data.output["distance_to_image_plane"][env_ids] = distance_to_image_plane.view( + wp.launch( + compute_distance_to_image_plane_masked_kernel, + dim=(self._num_envs, self.num_rays), + inputs=[ + env_mask, + self._quat_w_wp, + self._ray_distance, + self._ray_directions_w, + ], + outputs=[ + self._distance_to_image_plane_wp, + ], + device=self._device, + ) + # Apply depth clipping on the intermediate buffer (leaves _ray_distance unmodified) + self._apply_depth_clipping(env_mask, self._distance_to_image_plane_wp) + self._data.output["distance_to_image_plane"][env_ids] = self._distance_to_image_plane_torch[env_ids].view( -1, *self.image_shape, 1 ) if "distance_to_camera" in self.cfg.data_types: - if self.cfg.depth_clipping_behavior == "max": - ray_depth = torch.clip(ray_depth, max=self.cfg.max_distance) - elif self.cfg.depth_clipping_behavior == "zero": - ray_depth[ray_depth > self.cfg.max_distance] = 0.0 - self._data.output["distance_to_camera"][env_ids] = ray_depth.view(-1, *self.image_shape, 1) + # d2ip (if requested) was computed before this block so _ray_distance is still unclipped. + self._apply_depth_clipping(env_mask, self._ray_distance) + self._data.output["distance_to_camera"][env_ids] = self._ray_distance_torch[env_ids].view( + -1, *self.image_shape, 1 + ) if "normals" in self.cfg.data_types: - self._data.output["normals"][env_ids] = ray_normal.view(-1, *self.image_shape, 3) + self._data.output["normals"][env_ids] = self._ray_normal_w_torch[env_ids].view(-1, *self.image_shape, 3) def _debug_vis_callback(self, event): # in case it crashes be safe if not hasattr(self, "ray_hits_w"): return - # show ray hit positions - self.ray_visualizer.visualize(self.ray_hits_w.view(-1, 3)) + # filter out missed rays (inf values) before visualizing + ray_hits_flat = self.ray_hits_w.reshape(-1, 3) + valid_mask = ~torch.isinf(ray_hits_flat).any(dim=-1) + viz_points = ray_hits_flat[valid_mask] + # if no valid hits, skip + if viz_points.shape[0] == 0: + return + self.ray_visualizer.visualize(viz_points) """ Private Helpers """ + def _apply_depth_clipping(self, env_mask: wp.array, depth: wp.array) -> None: + """Apply depth clipping in-place on a warp float32 buffer. + + Uses :attr:`cfg.depth_clipping_behavior` to determine the fill value: + ``"max"`` replaces out-of-range and NaN values with :attr:`cfg.max_distance`; + ``"zero"`` replaces them with 0. No-op when behavior is ``"none"``. + + Args: + env_mask: Boolean mask selecting which environments to update. Shape is (num_envs,). + depth: Warp 2-D float32 buffer to clip in-place. Shape is (num_envs, num_rays). + """ + if self.cfg.depth_clipping_behavior == "max": + wp.launch( + apply_depth_clipping_masked_kernel, + dim=(self._num_envs, self.num_rays), + inputs=[env_mask, float(self.cfg.max_distance), float(self.cfg.max_distance), depth], + device=self._device, + ) + elif self.cfg.depth_clipping_behavior == "zero": + wp.launch( + apply_depth_clipping_masked_kernel, + dim=(self._num_envs, self.num_rays), + inputs=[env_mask, float(self.cfg.max_distance), float(0.0), depth], + device=self._device, + ) + elif self.cfg.depth_clipping_behavior == "none": + pass # no clipping: inf values remain as-is + else: + raise ValueError( + f"Unknown depth_clipping_behavior: {self.cfg.depth_clipping_behavior!r}." + " Valid values are 'max', 'zero', and 'none'." + ) + def _check_supported_data_types(self, cfg: RayCasterCameraCfg): """Checks if the data types are supported by the ray-caster camera.""" # check if there is any intersection in unsupported types @@ -362,7 +527,7 @@ def _check_supported_data_types(self, cfg: RayCasterCameraCfg): def _create_buffers(self): """Create buffers for storing data.""" - # prepare drift + # prepare drift (kept as torch tensors so subclasses may use torch indexing) self.drift = torch.zeros(self._view.count, 3, device=self.device) self.ray_cast_drift = torch.zeros(self._view.count, 3, device=self.device) # create the data object @@ -415,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. @@ -441,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] ) @@ -449,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_camera_cfg.py b/source/isaaclab/isaaclab/sensors/ray_caster/ray_caster_camera_cfg.py index 98020d845f82..574c95020437 100644 --- a/source/isaaclab/isaaclab/sensors/ray_caster/ray_caster_camera_cfg.py +++ b/source/isaaclab/isaaclab/sensors/ray_caster/ray_caster_camera_cfg.py @@ -54,8 +54,8 @@ class OffsetCfg: - ``"max"``: Values are clipped to the maximum value. - ``"zero"``: Values are clipped to zero. - - ``"none``: No clipping is applied. Values will be returned as ``inf`` for ``distance_to_camera`` and ``nan`` - for ``distance_to_image_plane`` data type. + - ``"none"``: No clipping is applied. Values will be returned as ``inf`` for missed rays in both + ``distance_to_camera`` and ``distance_to_image_plane`` data types. """ pattern_cfg: PinholeCameraPatternCfg = MISSING 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 7d91e446adac..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,30 +63,15 @@ 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". The options are: * ``base`` if the rays' starting positions and directions track the full root position and orientation. - * ``yaw`` if the rays' starting positions and directions track root position and only yaw component of - the orientation. This is useful for ray-casting height maps. + * ``yaw`` if the rays' starting positions track root position and the yaw component of the orientation, + while ray directions remain fixed in world frame. This is useful for ray-casting height maps where + the scan footprint should follow the body heading without tilting when the body pitches or rolls. * ``world`` if rays' starting positions and directions are always fixed. This is useful in combination with a mapping package on the robot and querying ray-casts in a global frame. """ diff --git a/source/isaaclab/isaaclab/sensors/ray_caster/ray_caster_data.py b/source/isaaclab/isaaclab/sensors/ray_caster/ray_caster_data.py index 6103a2167d66..c317c96f78b2 100644 --- a/source/isaaclab/isaaclab/sensors/ray_caster/ray_caster_data.py +++ b/source/isaaclab/isaaclab/sensors/ray_caster/ray_caster_data.py @@ -3,28 +3,75 @@ # # SPDX-License-Identifier: BSD-3-Clause -from dataclasses import dataclass +from __future__ import annotations -import torch +import warp as wp -@dataclass class RayCasterData: - """Data container for the ray-cast sensor.""" + """Data container for the ray-cast sensor. - pos_w: torch.Tensor = None - """Position of the sensor origin in world frame. - - Shape is (N, 3), where N is the number of sensors. + All public properties return :class:`wp.array` objects backed by device memory. + Use :func:`wp.to_torch` at the call-site when a PyTorch tensor is needed, e.g. + ``wp.to_torch(sensor.data.ray_hits_w)``. """ - quat_w: torch.Tensor = None - """Orientation of the sensor origin in quaternion (x, y, z, w) in world frame. - Shape is (N, 4), where N is the number of sensors. - """ - ray_hits_w: torch.Tensor = None - """The ray hit positions in the world frame. + def __init__(self): + self._pos_w: wp.array | None = None + self._quat_w: wp.array | None = None + self._ray_hits_w: wp.array | None = None - Shape is (N, B, 3), where N is the number of sensors, B is the number of rays - in the scan pattern per sensor. - """ + # Zero-copy torch views; kept alive to prevent GC of the underlying warp buffers. + # Not surfaced as public API — callers should use wp.to_torch() at the call-site. + self._pos_w_torch = None + self._quat_w_torch = None + self._ray_hits_w_torch = None + + @property + def pos_w(self) -> wp.array | None: + """Position of the sensor origin in world frame [m]. + + Shape is (N,), dtype ``wp.vec3f``. In torch this resolves to (N, 3), + where N is the number of sensors. Use :func:`wp.to_torch` to obtain a + :class:`torch.Tensor` view without copying data. + """ + return self._pos_w + + @property + def quat_w(self) -> wp.array | None: + """Orientation of the sensor origin in quaternion (x, y, z, w) in world frame. + + Shape is (N,), dtype ``wp.quatf``. In torch this resolves to (N, 4), + where N is the number of sensors. Use :func:`wp.to_torch` to obtain a + :class:`torch.Tensor` view without copying data. + """ + return self._quat_w + + @property + def ray_hits_w(self) -> wp.array | None: + """The ray hit positions in the world frame [m]. + + Shape is (N, B), dtype ``wp.vec3f``. In torch this resolves to (N, B, 3), + where N is the number of sensors and B is the number of rays per sensor. + Contains ``inf`` for missed hits. Use :func:`wp.to_torch` to obtain a + :class:`torch.Tensor` view without copying data. + """ + return self._ray_hits_w + + def create_buffers(self, num_envs: int, num_rays: int, device: str) -> None: + """Create internal warp buffers and corresponding zero-copy torch views. + + Args: + num_envs: Number of environments / sensors. + num_rays: Number of rays per sensor. + device: Device for tensor storage. + """ + self._device = device + + self._pos_w = wp.zeros(num_envs, dtype=wp.vec3f, device=device) + self._quat_w = wp.zeros(num_envs, dtype=wp.quatf, device=device) + self._ray_hits_w = wp.zeros((num_envs, num_rays), dtype=wp.vec3f, device=device) + + self._pos_w_torch = wp.to_torch(self._pos_w) + self._quat_w_torch = wp.to_torch(self._quat_w) + self._ray_hits_w_torch = wp.to_torch(self._ray_hits_w) 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/simulation_context.py b/source/isaaclab/isaaclab/sim/simulation_context.py index 5e05dab92a46..d63da8d60dd8 100644 --- a/source/isaaclab/isaaclab/sim/simulation_context.py +++ b/source/isaaclab/isaaclab/sim/simulation_context.py @@ -21,6 +21,7 @@ import isaaclab.sim as sim_utils import isaaclab.sim.utils.stage as stage_utils from isaaclab.app.settings_manager import SettingsManager +from isaaclab.envs.utils.recording_hooks import run_recording_hooks_after_visualizers from isaaclab.physics import BaseSceneDataProvider, PhysicsManager, SceneDataProvider from isaaclab.physics.scene_data_requirements import ( SceneDataRequirement, @@ -28,6 +29,7 @@ resolve_scene_data_requirements, ) from isaaclab.sim.utils import create_new_stage +from isaaclab.utils.string import clear_resolve_matching_names_cache from isaaclab.utils.version import has_kit from isaaclab.visualizers.base_visualizer import BaseVisualizer @@ -181,6 +183,7 @@ def __init__(self, cfg: SimulationCfg | None = None): self._has_offscreen_render = bool(self.get_setting("/isaaclab/render/offscreen")) self._xr_enabled = bool(self.get_setting("/isaaclab/xr/enabled")) # Note: has_rtx_sensors is NOT cached because it changes when Camera sensors are created + self._pending_camera_view: tuple[tuple[float, float, float], tuple[float, float, float]] | None = None # Simulation state self._is_playing = False @@ -188,6 +191,10 @@ def __init__(self, cfg: SimulationCfg | None = None): # Monotonic physics-step counter used by camera sensors for self._physics_step_count: int = 0 + # Monotonic render-generation counter. This increments whenever render() + # is executed and lets downstream camera freshness logic distinguish + # render/reset transitions that occur without advancing physics steps. + self._render_generation: int = 0 type(self)._instance = self # Mark as valid singleton only after successful init @@ -290,7 +297,8 @@ def _init_usd_physics_scene(self) -> None: UsdPhysics.SetStageKilogramsPerUnit(self.stage, 1.0) # Find and delete any existing physics scene. - # Collect paths first to avoid iterator invalidation during deletion. + # Collect paths first to avoid mutating the stage while traversing, + # which can invalidate the USD iterator. physics_scene_paths = [ prim.GetPath().pathString for prim in self.stage.Traverse() if prim.GetTypeName() == "PhysicsScene" ] @@ -340,6 +348,16 @@ def has_offscreen_render(self) -> bool: """Returns whether offscreen rendering is enabled (cached at init).""" return self._has_offscreen_render + def has_active_visualizers(self) -> bool: + """Return whether any visualizer path is active for rendering/camera control.""" + return bool(self.get_setting("/isaaclab/visualizer/types")) or bool( + self.get_setting("/isaaclab/video/auto_start_kit") + ) + + def can_render_rgb_array(self) -> bool: + """Return whether rgb-array rendering is currently available.""" + return self.has_gui or self.has_offscreen_render or self.has_active_visualizers() + @property def is_rendering(self) -> bool: """Returns whether rendering is active (GUI, RTX sensors, visualizers, or XR).""" @@ -355,6 +373,11 @@ def get_physics_dt(self) -> float: """Returns the physics time step.""" return self.physics_manager.get_physics_dt() + @property + def render_generation(self) -> int: + """Returns a monotonic counter for render() executions.""" + return self._render_generation + def _create_default_visualizer_configs(self, requested_visualizers: list[str]) -> list: """Create default visualizer configs for requested types. @@ -408,39 +431,23 @@ def _get_cli_visualizer_types(self) -> list[str]: # App launcher writes this as a single string; accept comma and/or whitespace separators. return [value for chunk in requested.split(",") for value in chunk.split() if value] - def _get_cli_visualizer_max_worlds_override(self) -> tuple[bool, int | None]: - """Return CLI override for visualizer max worlds. + def _apply_visualizer_cli_overrides(self, visualizer_cfgs: list[Any]) -> None: + """Apply ``--max_visible_envs`` to every resolved visualizer cfg when set in settings. - Returns: - Tuple of (has_override, value), where value=None means no override. + AppLauncher stores ``/isaaclab/visualizer/max_visible_envs`` as ``-1`` when the flag was + omitted; any non-negative int overrides :attr:`VisualizerCfg.max_visible_envs` on each cfg. """ - value = self.get_setting("/isaaclab/visualizer/max_worlds") - if value is None: - return False, None + raw = self.get_setting("/isaaclab/visualizer/max_visible_envs") try: - max_worlds = int(value) + max_visible = int(raw) if raw is not None else -1 except (TypeError, ValueError): - logger.warning("[SimulationContext] Invalid /isaaclab/visualizer/max_worlds setting: %r", value) - return False, None - - # -1 means no CLI override. - if max_worlds < 0: - return False, None - return True, max_worlds - - def _apply_visualizer_cli_overrides(self, visualizer_cfgs: list[Any]) -> None: - """Apply CLI visualizer overrides (e.g., max worlds) to resolved configs. - - Args: - visualizer_cfgs: Resolved visualizer configs to update in-place. - """ - has_max_worlds_override, max_worlds_override = self._get_cli_visualizer_max_worlds_override() - if not has_max_worlds_override: + logger.warning("[SimulationContext] Invalid /isaaclab/visualizer/max_visible_envs: %r", raw) + return + if max_visible < 0: return - for cfg in visualizer_cfgs: - if hasattr(cfg, "max_worlds"): - cfg.max_worlds = max_worlds_override + if hasattr(cfg, "max_visible_envs"): + cfg.max_visible_envs = max_visible def _is_cli_visualizer_explicit(self) -> bool: """Return ``True`` when visualizers were explicitly provided via CLI.""" @@ -580,6 +587,14 @@ def initialize_visualizers(self) -> None: exc, ) + # Replay any camera pose requested before visualizers were initialized. + pending = getattr(self, "_pending_camera_view", None) + if pending is not None: + eye, target = pending + for viz in self._visualizers: + viz.set_camera_view(eye, target) + self._pending_camera_view = None + if not self._visualizers and self._scene_data_provider is not None: close_provider = getattr(self._scene_data_provider, "close", None) if callable(close_provider): @@ -630,6 +645,7 @@ def get_rendering_dt(self) -> float: def set_camera_view(self, eye: tuple, target: tuple) -> None: """Set camera view on all visualizers that support it.""" + self._pending_camera_view = (tuple(eye), tuple(target)) for viz in self._visualizers: viz.set_camera_view(eye, target) @@ -676,10 +692,15 @@ def render(self, mode: int | None = None) -> None: Calls update_visualizers() so visualizers run at the render cadence (not at every physics step). Camera sensors drive their configured renderer when - fetching data, so this method remains backend-agnostic. + fetching data. Recording-related follow-up (Kit/RTX headless video, Newton GL + video, etc.) runs in :mod:`isaaclab.envs.utils.recording_hooks` so it is not tied to a + specific :class:`~isaaclab.physics.PhysicsManager` subclass. """ self.physics_manager.pre_render() self.update_visualizers(self.get_rendering_dt()) + self.physics_manager.after_visualizers_render() + run_recording_hooks_after_visualizers(self) + self._render_generation += 1 # Call render callbacks if hasattr(self, "_render_callbacks"): @@ -704,6 +725,11 @@ def update_visualizers(self, dt: float) -> None: visualizers_to_remove.append(viz) continue if viz.is_rendering_paused(): + # Keep non-Kit visualizer event loops responsive while rendering is paused. + # Newton/Rerun/Viser need step(0.0) so GL/UI can process input (e.g. Resume). + # Kit is skipped: step() would call app.update(), which must not run during pause. + if not viz.pumps_app_update(): + viz.step(0.0) continue while viz.is_training_paused() and viz.is_running(): viz.step(0.0) @@ -726,14 +752,7 @@ def update_scene_data_provider(self, force_require_forward: bool = False): self._visualizer_step_counter += 1 if self._scene_data_provider is None: return - provider = self._scene_data_provider - env_ids_union: list[int] = [] - for viz in self._visualizers: - ids = viz.get_visualized_env_ids() - if ids is not None: - env_ids_union.extend(ids) - env_ids = list(dict.fromkeys(env_ids_union)) if env_ids_union else None - provider.update(env_ids) + self._scene_data_provider.update() def _should_forward_before_visualizer_update(self) -> bool: """Return True if any visualizer requires pre-step forward kinematics.""" @@ -800,6 +819,9 @@ def clear_instance(cls) -> None: # close_stage() + app shutdown destroy the entire stage at once. stage_utils.close_stage() + # Discard cached name-resolution data from destroyed assets + clear_resolve_matching_names_cache() + # Clear instance cls._instance = None 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/isaaclab/test/mock_interfaces/utils/mock_wrench_composer.py b/source/isaaclab/isaaclab/test/mock_interfaces/utils/mock_wrench_composer.py index f35228ea6dcc..a20dbeb01274 100644 --- a/source/isaaclab/isaaclab/test/mock_interfaces/utils/mock_wrench_composer.py +++ b/source/isaaclab/isaaclab/test/mock_interfaces/utils/mock_wrench_composer.py @@ -11,6 +11,7 @@ from __future__ import annotations +import warnings from typing import TYPE_CHECKING import torch @@ -21,14 +22,16 @@ class MockWrenchComposer: - """Mock WrenchComposer for testing. + """Mock WrenchComposer matching the dual-buffer API for testing. This class provides a mock implementation of WrenchComposer that matches the real interface but does not launch Warp kernels. It can be used for testing and benchmarking asset classes without requiring the full simulation environment. - The mock maintains simple buffers and sets the active flag when forces/torques are added, - but does not perform actual force composition computations. + The mock maintains the 5 input buffers and 2 output buffers matching the real WrenchComposer, + and sets the active flag when forces/torques are added. The ``compose_to_body_frame()`` method + simply copies the local buffers to the output buffers (since mock assets typically use identity + transforms). """ def __init__(self, asset: BaseArticulation | BaseRigidObject | BaseRigidObjectCollection) -> None: @@ -44,15 +47,23 @@ def __init__(self, asset: BaseArticulation | BaseRigidObject | BaseRigidObjectCo raise ValueError(f"Unsupported asset type: {asset.__class__.__name__}") self.device = asset.device self._asset = asset - self._active = False - # Create buffers using Warp (matching real WrenchComposer) - self._composed_force_b = wp.zeros((self.num_envs, self.num_bodies), dtype=wp.vec3f, device=self.device) - self._composed_torque_b = wp.zeros((self.num_envs, self.num_bodies), dtype=wp.vec3f, device=self.device) + # -- Tracking flags -- + self._active: bool = False + self._dirty: bool = False - # Create torch views (matching real WrenchComposer) - self._composed_force_b_torch = wp.to_torch(self._composed_force_b) - self._composed_torque_b_torch = wp.to_torch(self._composed_torque_b) + shape = (self.num_envs, self.num_bodies) + + # -- 5 input buffers -- + self._global_force_w = wp.zeros(shape, dtype=wp.vec3f, device=self.device) + self._global_torque_w = wp.zeros(shape, dtype=wp.vec3f, device=self.device) + self._global_force_at_com_w = wp.zeros(shape, dtype=wp.vec3f, device=self.device) + self._local_force_b = wp.zeros(shape, dtype=wp.vec3f, device=self.device) + self._local_torque_b = wp.zeros(shape, dtype=wp.vec3f, device=self.device) + + # -- 2 output buffers -- + self._out_force_b = wp.zeros(shape, dtype=wp.vec3f, device=self.device) + self._out_torque_b = wp.zeros(shape, dtype=wp.vec3f, device=self.device) # Create index arrays self._ALL_ENV_INDICES_WP = wp.from_torch( @@ -64,46 +75,145 @@ def __init__(self, asset: BaseArticulation | BaseRigidObject | BaseRigidObjectCo self._ALL_ENV_INDICES_TORCH = wp.to_torch(self._ALL_ENV_INDICES_WP) self._ALL_BODY_INDICES_TORCH = wp.to_torch(self._ALL_BODY_INDICES_WP) + # ------------------------------------------------------------------ + # Properties + # ------------------------------------------------------------------ + @property def active(self) -> bool: - """Whether the wrench composer is active.""" + """Whether any forces or torques have been written since the last full reset.""" return self._active + # -- Input buffer accessors (read-only) -- + + @property + def global_force_w(self) -> wp.array: + """Positional global forces buffer. Shape ``(num_envs, num_bodies)``, dtype ``wp.vec3f``.""" + return self._global_force_w + + @property + def global_torque_w(self) -> wp.array: + """Global torques buffer (about world origin). Shape ``(num_envs, num_bodies)``, dtype ``wp.vec3f``.""" + return self._global_torque_w + + @property + def global_force_at_com_w(self) -> wp.array: + """Global forces at CoM buffer (no positional torque). Shape ``(num_envs, num_bodies)``, dtype ``wp.vec3f``.""" + return self._global_force_at_com_w + + @property + def local_force_b(self) -> wp.array: + """Body-frame forces buffer. Shape ``(num_envs, num_bodies)``, dtype ``wp.vec3f``.""" + return self._local_force_b + + @property + def local_torque_b(self) -> wp.array: + """Body-frame torques buffer. Shape ``(num_envs, num_bodies)``, dtype ``wp.vec3f``.""" + return self._local_torque_b + + # -- Output buffer accessors -- + + @property + def out_force_b(self) -> wp.array: + """Composed force in the body (link) frame. Shape ``(num_envs, num_bodies)``, dtype ``wp.vec3f``. + + Triggers composition from input buffers if dirty. + """ + self._ensure_composed() + return self._out_force_b + + @property + def out_torque_b(self) -> wp.array: + """Composed torque in the body (link) frame. Shape ``(num_envs, num_bodies)``, dtype ``wp.vec3f``. + + Triggers composition from input buffers if dirty. + """ + self._ensure_composed() + return self._out_torque_b + + # -- Legacy composed_force / composed_torque properties for backward compat -- + @property def composed_force(self) -> wp.array: """Composed force at the body's link frame. - Returns: - wp.array: Composed force at the body's link frame. (num_envs, num_bodies, 3) + .. deprecated:: 4.5.33 + Use :attr:`out_force_b` instead. """ - return self._composed_force_b + warnings.warn( + "The property 'composed_force' is deprecated. Use 'out_force_b' instead.", + DeprecationWarning, + stacklevel=2, + ) + return self.out_force_b @property def composed_torque(self) -> wp.array: """Composed torque at the body's link frame. - Returns: - wp.array: Composed torque at the body's link frame. (num_envs, num_bodies, 3) + .. deprecated:: 4.5.33 + Use :attr:`out_torque_b` instead. """ - return self._composed_torque_b + warnings.warn( + "The property 'composed_torque' is deprecated. Use 'out_torque_b' instead.", + DeprecationWarning, + stacklevel=2, + ) + return self.out_torque_b - @property - def composed_force_as_torch(self) -> torch.Tensor: - """Composed force at the body's link frame as torch tensor. + # ------------------------------------------------------------------ + # Composition + # ------------------------------------------------------------------ + + def compose_to_body_frame(self): + """Mock composition: sums all input buffers to output assuming identity transforms. - Returns: - torch.Tensor: Composed force at the body's link frame. (num_envs, num_bodies, 3) + Under identity transforms (no rotation), global-frame values equal body-frame values, + so all five input buffers are summed directly into the two output buffers. """ - return self._composed_force_b_torch + # Zero output buffers + self._out_force_b.zero_() + self._out_torque_b.zero_() - @property - def composed_torque_as_torch(self) -> torch.Tensor: - """Composed torque at the body's link frame as torch tensor. + # Use torch views for the accumulation + out_force_torch = wp.to_torch(self._out_force_b) + out_torque_torch = wp.to_torch(self._out_torque_b) + + # Sum all force contributions (identity: no rotation needed) + out_force_torch.add_(wp.to_torch(self._local_force_b)) + out_force_torch.add_(wp.to_torch(self._global_force_w)) + out_force_torch.add_(wp.to_torch(self._global_force_at_com_w)) - Returns: - torch.Tensor: Composed torque at the body's link frame. (num_envs, num_bodies, 3) + # Sum all torque contributions + out_torque_torch.add_(wp.to_torch(self._local_torque_b)) + out_torque_torch.add_(wp.to_torch(self._global_torque_w)) + + self._dirty = False + + # ------------------------------------------------------------------ + # Buffer merging + # ------------------------------------------------------------------ + + def add_raw_buffers_from(self, other: MockWrenchComposer): + """Element-wise add another composer's five input buffers into this one. + + Args: + other: Another :class:`MockWrenchComposer` whose input buffers will be added into this one. """ - return self._composed_torque_b_torch + # Use torch views for element-wise addition + wp.to_torch(self._global_force_w).add_(wp.to_torch(other._global_force_w)) + wp.to_torch(self._global_torque_w).add_(wp.to_torch(other._global_torque_w)) + wp.to_torch(self._global_force_at_com_w).add_(wp.to_torch(other._global_force_at_com_w)) + wp.to_torch(self._local_force_b).add_(wp.to_torch(other._local_force_b)) + wp.to_torch(self._local_torque_b).add_(wp.to_torch(other._local_torque_b)) + + if other._active: + self._active = True + self._dirty = True + + # ------------------------------------------------------------------ + # Add / Set methods + # ------------------------------------------------------------------ def add_forces_and_torques( self, @@ -172,9 +282,10 @@ def add_forces_and_torques_index( env_ids: torch.Tensor | None = None, is_global: bool = False, ) -> None: - """Add forces and torques by index (mock - just sets active flag).""" + """Add forces and torques by index (mock - sets active/dirty flags).""" if forces is not None or torques is not None: self._active = True + self._dirty = True def add_forces_and_torques_mask( self, @@ -185,9 +296,10 @@ def add_forces_and_torques_mask( env_mask: wp.array | torch.Tensor | None = None, is_global: bool = False, ) -> None: - """Add forces and torques by mask (mock - just sets active flag).""" + """Add forces and torques by mask (mock - sets active/dirty flags).""" if forces is not None or torques is not None: self._active = True + self._dirty = True def set_forces_and_torques_index( self, @@ -198,9 +310,10 @@ def set_forces_and_torques_index( env_ids: wp.array | torch.Tensor | None = None, is_global: bool = False, ) -> None: - """Set forces and torques by index (mock - just sets active flag).""" + """Set forces and torques by index (mock - sets active/dirty flags).""" if forces is not None or torques is not None: self._active = True + self._dirty = True def set_forces_and_torques_mask( self, @@ -211,28 +324,65 @@ def set_forces_and_torques_mask( env_mask: wp.array | torch.Tensor | None = None, is_global: bool = False, ) -> None: - """Set forces and torques by mask (mock - just sets active flag).""" + """Set forces and torques by mask (mock - sets active/dirty flags).""" if forces is not None or torques is not None: self._active = True + self._dirty = True + + # ------------------------------------------------------------------ + # Reset + # ------------------------------------------------------------------ def reset(self, env_ids: wp.array | torch.Tensor | None = None, env_mask: wp.array | None = None) -> None: - """Reset the composed force and torque. + """Reset all 7 buffers (5 input + 2 output) and clear all flags. Args: env_ids: Environment ids to reset. Defaults to None (all environments). env_mask: Environment mask to reset. Defaults to None (all environments). """ - if env_ids is None: - self._composed_force_b.zero_() - self._composed_torque_b.zero_() + if env_ids is None and env_mask is None: + # Full reset: zero all 7 buffers and clear flags + self._global_force_w.zero_() + self._global_torque_w.zero_() + self._global_force_at_com_w.zero_() + self._local_force_b.zero_() + self._local_torque_b.zero_() + self._out_force_b.zero_() + self._out_torque_b.zero_() self._active = False + self._dirty = False else: - # For partial reset, just zero the specified environments + # For partial reset, just zero the specified environments across all 7 buffers if isinstance(env_ids, torch.Tensor): indices = wp.from_torch(env_ids.to(torch.int32), dtype=wp.int32) elif isinstance(env_ids, list): indices = wp.array(env_ids, dtype=wp.int32, device=self.device) else: indices = env_ids - self._composed_force_b[indices].zero_() - self._composed_torque_b[indices].zero_() + + # Zero all 7 buffers for the specified environments + # Use torch views for the indexing operation + for buf in [ + self._global_force_w, + self._global_torque_w, + self._global_force_at_com_w, + self._local_force_b, + self._local_torque_b, + self._out_force_b, + self._out_torque_b, + ]: + buf_torch = wp.to_torch(buf) + if isinstance(env_ids, torch.Tensor): + buf_torch[env_ids.long()] = 0.0 + else: + idx_torch = wp.to_torch(indices).long() + buf_torch[idx_torch] = 0.0 + + # ------------------------------------------------------------------ + # Internal helpers + # ------------------------------------------------------------------ + + def _ensure_composed(self): + """Compose input buffers into output buffers if dirty.""" + if self._dirty: + self.compose_to_body_frame() diff --git a/source/isaaclab/isaaclab/utils/__init__.pyi b/source/isaaclab/isaaclab/utils/__init__.pyi index 84d6e8b7f098..1ca7ef7866c6 100644 --- a/source/isaaclab/isaaclab/utils/__init__.pyi +++ b/source/isaaclab/isaaclab/utils/__init__.pyi @@ -46,6 +46,7 @@ __all__ = [ "string_to_callable", "ResolvableString", "resolve_matching_names", + "clear_resolve_matching_names_cache", "resolve_matching_names_values", "find_unique_string_name", "find_root_prim_path_from_regex", @@ -98,6 +99,7 @@ from .string import ( string_to_callable, ResolvableString, resolve_matching_names, + clear_resolve_matching_names_cache, resolve_matching_names_values, find_unique_string_name, find_root_prim_path_from_regex, diff --git a/source/isaaclab/isaaclab/utils/string.py b/source/isaaclab/isaaclab/utils/string.py index c4033055d8a3..4e7790006ade 100644 --- a/source/isaaclab/isaaclab/utils/string.py +++ b/source/isaaclab/isaaclab/utils/string.py @@ -6,6 +6,7 @@ """Sub-module containing utilities for transforming strings and regular expressions.""" import ast +import functools import importlib import inspect import re @@ -247,50 +248,19 @@ def __deepcopy__(self, memo): """ -def resolve_matching_names( - keys: str | Sequence[str], - list_of_strings: Sequence[str], - preserve_order: bool = False, - *, - raise_when_no_match: bool = True, -) -> tuple[list[int], list[str]]: - """Match a list of query regular expressions against a list of strings and return the matched indices and names. - - When a list of query regular expressions is provided, the function checks each target string against each - query regular expression and returns the indices of the matched strings and the matched strings. - - If the :attr:`preserve_order` is True, the ordering of the matched indices and names is the same as the order - of the provided list of strings. This means that the ordering is dictated by the order of the target strings - and not the order of the query regular expressions. - - If the :attr:`preserve_order` is False, the ordering of the matched indices and names is the same as the order - of the provided list of query regular expressions. - - For example, consider the list of strings is ['a', 'b', 'c', 'd', 'e'] and the regular expressions are ['a|c', 'b']. - If :attr:`preserve_order` is False, then the function will return the indices of the matched strings and the - strings as: ([0, 1, 2], ['a', 'b', 'c']). When :attr:`preserve_order` is True, it will return them as: - ([0, 2, 1], ['a', 'c', 'b']). +@functools.cache +def _resolve_matching_names_impl( + keys: tuple[str, ...], + list_of_strings: tuple[str, ...], + preserve_order: bool, + raise_when_no_match: bool, +) -> tuple[tuple[int, ...], tuple[str, ...]]: + """Cached implementation of :func:`resolve_matching_names`. - Note: - The function does not sort the indices. It returns the indices in the order they are found. - - Args: - keys: A regular expression or a list of regular expressions to match the strings in the list. - list_of_strings: A list of strings to match. - preserve_order: Whether to preserve the order of the query keys in the returned values. Defaults to False. - raise_when_no_match: Whether to raise a ``ValueError`` when not all regular expressions are matched. - Defaults to True. When False, returns empty lists instead of raising. - - Returns: - A tuple of lists containing the matched indices and names. - - Raises: - ValueError: When multiple matches are found for a string in the list. - ValueError: When not all regular expressions are matched and :attr:`raise_when_no_match` is True. + All arguments are hashable so that ``functools.cache`` can store results. + Returns tuples (immutable) to protect the cached data from mutation; + the public wrapper converts these back to fresh lists for each caller. """ - # resolve name keys - if isinstance(keys, str): - keys = [keys] # find matching patterns index_list = [] names_list = [] @@ -337,7 +307,7 @@ def resolve_matching_names( # check that all regular expressions are matched if not all(keys_match_found): if not raise_when_no_match: - return [], [] + return (), () # make this print nicely aligned for debugging msg = "\n" for key, value in zip(keys, keys_match_found): @@ -347,8 +317,66 @@ def resolve_matching_names( raise ValueError( f"Not all regular expressions are matched! Please check that the regular expressions are correct: {msg}" ) - # return - return index_list, names_list + # return immutable tuples for safe caching + return tuple(index_list), tuple(names_list) + + +def resolve_matching_names( + keys: str | Sequence[str], + list_of_strings: Sequence[str], + preserve_order: bool = False, + *, + raise_when_no_match: bool = True, +) -> tuple[list[int], list[str]]: + """Match a list of query regular expressions against a list of strings and return the matched indices and names. + + When a list of query regular expressions is provided, the function checks each target string against each + query regular expression and returns the indices of the matched strings and the matched strings. + + If the :attr:`preserve_order` is True, the ordering of the matched indices and names is the same as the order + of the provided list of strings. This means that the ordering is dictated by the order of the target strings + and not the order of the query regular expressions. + + If the :attr:`preserve_order` is False, the ordering of the matched indices and names is the same as the order + of the provided list of query regular expressions. + + For example, consider the list of strings is ['a', 'b', 'c', 'd', 'e'] and the regular expressions are ['a|c', 'b']. + If :attr:`preserve_order` is False, then the function will return the indices of the matched strings and the + strings as: ([0, 1, 2], ['a', 'b', 'c']). When :attr:`preserve_order` is True, it will return them as: + ([0, 2, 1], ['a', 'c', 'b']). + + Results are cached internally — repeated calls with the same arguments avoid redundant regex matching. + + Note: + The function does not sort the indices. It returns the indices in the order they are found. + + Args: + keys: A regular expression or a list of regular expressions to match the strings in the list. + list_of_strings: A list of strings to match. + preserve_order: Whether to preserve the order of the query keys in the returned values. Defaults to False. + raise_when_no_match: Whether to raise a ``ValueError`` when not all regular expressions are matched. + Defaults to True. When False, returns empty lists instead of raising. + + Returns: + A tuple of lists containing the matched indices and names. + + Raises: + ValueError: When multiple matches are found for a string in the list. + ValueError: When not all regular expressions are matched and :attr:`raise_when_no_match` is True. + """ + _keys = (keys,) if isinstance(keys, str) else tuple(keys) + idx, names = _resolve_matching_names_impl(_keys, tuple(list_of_strings), preserve_order, raise_when_no_match) + return list(idx), list(names) + + +def clear_resolve_matching_names_cache() -> None: + """Discard all cached results from :func:`resolve_matching_names`. + + Call this when the simulation scene is torn down so that cached + name-resolution entries from destroyed assets do not accumulate + across scene rebuilds in long-lived processes. + """ + _resolve_matching_names_impl.cache_clear() def resolve_matching_names_values( @@ -360,6 +388,11 @@ def resolve_matching_names_values( """Match a list of regular expressions in a dictionary against a list of strings and return the matched indices, names, and values. + Note: + Unlike :func:`resolve_matching_names`, this function is not cached. Current callers + use it during initialization only (e.g. action/actuator config resolution), so caching + would add complexity without a measurable benefit. + If the :attr:`preserve_order` is True, the ordering of the matched indices and names is the same as the order of the provided list of strings. This means that the ordering is dictated by the order of the target strings and not the order of the query regular expressions. diff --git a/source/isaaclab/isaaclab/utils/warp/kernels.py b/source/isaaclab/isaaclab/utils/warp/kernels.py index da2d9123db47..efcdbfe63f1e 100644 --- a/source/isaaclab/isaaclab/utils/warp/kernels.py +++ b/source/isaaclab/isaaclab/utils/warp/kernels.py @@ -53,7 +53,7 @@ def raycast_mesh_kernel( this array is not used. max_dist: The maximum ray-cast distance. Defaults to 1e6. return_distance: Whether to return the ray hit distances. Defaults to False. - return_normal: Whether to return the ray hit normals. Defaults to False`. + return_normal: Whether to return the ray hit normals. Defaults to False. return_face_id: Whether to return the ray hit face ids. Defaults to False. """ # get the thread id @@ -79,6 +79,63 @@ def raycast_mesh_kernel( ray_face_id[tid] = f +@wp.kernel(enable_backward=False) +def raycast_mesh_masked_kernel( + # input + mesh: wp.uint64, + env_mask: wp.array(dtype=wp.bool), + ray_starts: wp.array2d(dtype=wp.vec3f), + ray_directions: wp.array2d(dtype=wp.vec3f), + max_dist: wp.float32, + return_distance: int, + return_normal: int, + # output + ray_hits: wp.array2d(dtype=wp.vec3f), + ray_distance: wp.array2d(dtype=wp.float32), + ray_normal: wp.array2d(dtype=wp.vec3f), +): + """Ray-cast against a single static mesh for masked environments. + + Extends :func:`raycast_mesh_kernel` with environment masking and optional distance/normal output, + for use in multi-environment sensor pipelines. + + Launch with ``dim=(num_envs, num_rays)``. + + Args: + mesh: Warp mesh id to ray-cast against. + env_mask: Boolean mask for which environments to update. Shape is (num_envs,). + ray_starts: World-frame ray start positions [m]. Shape is (num_envs, num_rays). + ray_directions: World-frame unit ray directions. Shape is (num_envs, num_rays). + max_dist: Maximum ray-cast distance [m]. + return_distance: Whether to write hit distances to ``ray_distance`` (1) or skip (0). + return_normal: Whether to write surface normals to ``ray_normal`` (1) or skip (0). + ray_hits: Output ray hit positions [m]. Shape is (num_envs, num_rays). + Pre-filled with inf for missed hits; unchanged on miss. + ray_distance: Output hit distances [m]. Shape is (num_envs, num_rays). + Written only when ``return_distance`` is 1; pre-filled with inf for missed hits. + ray_normal: Output surface normals at hit positions. Shape is (num_envs, num_rays). + Written only when ``return_normal`` is 1; pre-filled with inf for missed hits. + """ + env, ray = wp.tid() + if not env_mask[env]: + return + + t = float(0.0) + u = float(0.0) + v = float(0.0) + sign = float(0.0) + n = wp.vec3f() + f = int(0) + + hit = wp.mesh_query_ray(mesh, ray_starts[env, ray], ray_directions[env, ray], max_dist, t, u, v, sign, n, f) + if hit: + ray_hits[env, ray] = ray_starts[env, ray] + t * ray_directions[env, ray] + if return_distance == 1: + ray_distance[env, ray] = t + if return_normal == 1: + ray_normal[env, ray] = n + + @wp.kernel(enable_backward=False) def raycast_static_meshes_kernel( mesh: wp.array2d(dtype=wp.uint64), @@ -110,14 +167,24 @@ def raycast_static_meshes_kernel( account the mesh's position and rotation. This kernel is useful for ray-casting against static meshes that are not expected to move. + .. warning:: + **Known race condition:** When two meshes are equidistant to the same ray, the + ``atomic_min`` + equality-check pattern used for closest-hit resolution is not fully + thread-safe. Two threads may both pass the equality check and write different output + fields (e.g., ``ray_hits`` from mesh A, ``ray_normal`` from mesh B). In practice this + is rare (requires exact floating-point tie) and the position output is still correct, + but normals, face IDs, and mesh IDs may be inconsistent for the affected ray. + See `warp#1058 `_ for progress on a + thread-safe fix. + Args: mesh: The input mesh. The ray-casting is performed against this mesh on the device specified by the `mesh`'s `device` attribute. ray_starts: The input ray start positions. Shape is (B, N, 3). ray_directions: The input ray directions. Shape is (B, N, 3). ray_hits: The output ray hit positions. Shape is (B, N, 3). - ray_distance: The output ray hit distances. Shape is (B, N,), if ``return_distance`` is True. Otherwise, - this array is not used. + ray_distance: The closest hit distance buffer. Shape is (B, N). Updated via ``atomic_min`` for every + thread that records a hit; used to resolve closest-hit among multiple meshes. ray_normal: The output ray hit normals. Shape is (B, N, 3), if ``return_normal`` is True. Otherwise, this array is not used. ray_face_id: The output ray hit face ids. Shape is (B, N,), if ``return_face_id`` is True. Otherwise, @@ -125,7 +192,7 @@ def raycast_static_meshes_kernel( ray_mesh_id: The output ray hit mesh ids. Shape is (B, N,), if ``return_mesh_id`` is True. Otherwise, this array is not used. max_dist: The maximum ray-cast distance. Defaults to 1e6. - return_normal: Whether to return the ray hit normals. Defaults to False`. + return_normal: Whether to return the ray hit normals. Defaults to False. return_face_id: Whether to return the ray hit face ids. Defaults to False. return_mesh_id: Whether to return the mesh id. Defaults to False. """ @@ -141,10 +208,11 @@ def raycast_static_meshes_kernel( # if the ray hit, store the hit data if mesh_query_ray_t.result: wp.atomic_min(ray_distance, tid_env, tid_ray, mesh_query_ray_t.t) - # check if hit distance is less than the current hit distance, only then update the memory - # TODO, in theory we could use the output of atomic_min to avoid the non-thread safe next comparison - # however, warp atomic_min is returning the wrong values on gpu currently. - # FIXME https://github.com/NVIDIA/warp/issues/1058 + # TODO(warp#1058): Use the return value of atomic_min to avoid the non-thread-safe + # equality check below. Currently warp atomic_min returns wrong values on GPU, so we + # fall back to a racy read-back. When two meshes tie on distance, normals/face-ids/ + # mesh-ids may be written by different threads. The hit *position* is still correct + # because all tying threads compute the same world-space point. if mesh_query_ray_t.t == ray_distance[tid_env, tid_ray]: # convert back to world space and update the hit data ray_hits[tid_env, tid_ray] = start_pos + mesh_query_ray_t.t * direction @@ -160,6 +228,7 @@ def raycast_static_meshes_kernel( @wp.kernel(enable_backward=False) def raycast_dynamic_meshes_kernel( + env_mask: wp.array(dtype=wp.bool), mesh: wp.array2d(dtype=wp.uint64), ray_starts: wp.array2d(dtype=wp.vec3), ray_directions: wp.array2d(dtype=wp.vec3), @@ -175,7 +244,7 @@ def raycast_dynamic_meshes_kernel( return_face_id: int = False, return_mesh_id: int = False, ): - """Performs ray-casting against multiple meshes. + """Performs ray-casting against multiple dynamic meshes. This function performs ray-casting against the given meshes using the provided ray start positions and directions. The resulting ray hit positions are stored in the :obj:`ray_hits` array. @@ -183,7 +252,6 @@ def raycast_dynamic_meshes_kernel( The function utilizes the ``mesh_query_ray`` method from the ``wp`` module to perform the actual ray-casting operation. The maximum ray-cast distance is set to ``1e6`` units. - Note: That the ``ray_starts``, ``ray_directions``, and ``ray_hits`` arrays should have compatible shapes and data types to ensure proper execution. Additionally, they all must be in the same frame. @@ -191,29 +259,42 @@ def raycast_dynamic_meshes_kernel( All arguments are expected to be batched with the first dimension (B, batch) being the number of envs and the second dimension (N, num_rays) being the number of rays. For Meshes, W is the number of meshes. + .. warning:: + **Known race condition:** When two meshes are equidistant to the same ray, the + ``atomic_min`` + equality-check pattern used for closest-hit resolution is not fully + thread-safe. Two threads may both pass the equality check and write different output + fields (e.g., ``ray_hits`` from mesh A, ``ray_normal`` from mesh B). In practice this + is rare (requires exact floating-point tie) and the position output is still correct, + but normals, face IDs, and mesh IDs may be inconsistent for the affected ray. + See `warp#1058 `_ for progress on a + thread-safe fix. + Args: + env_mask: Boolean mask selecting which environments to process. Shape is (B,). mesh: The input mesh. The ray-casting is performed against this mesh on the device specified by the `mesh`'s `device` attribute. ray_starts: The input ray start positions. Shape is (B, N, 3). ray_directions: The input ray directions. Shape is (B, N, 3). ray_hits: The output ray hit positions. Shape is (B, N, 3). - ray_distance: The output ray hit distances. Shape is (B, N,), if ``return_distance`` is True. Otherwise, - this array is not used. + ray_distance: The closest hit distance buffer. Shape is (B, N). Updated via ``atomic_min`` for every + thread that records a hit; used to resolve closest-hit among multiple meshes. ray_normal: The output ray hit normals. Shape is (B, N, 3), if ``return_normal`` is True. Otherwise, this array is not used. ray_face_id: The output ray hit face ids. Shape is (B, N,), if ``return_face_id`` is True. Otherwise, this array is not used. ray_mesh_id: The output ray hit mesh ids. Shape is (B, N,), if ``return_mesh_id`` is True. Otherwise, this array is not used. - mesh_positions: The input mesh positions in world frame. Shape is (W, 3). - mesh_rotations: The input mesh rotations in world frame. Shape is (W, 4). + mesh_positions: The input mesh positions in world frame. Shape is (B, W, 3). + mesh_rotations: The input mesh rotations in world frame. Shape is (B, W, 4). max_dist: The maximum ray-cast distance. Defaults to 1e6. - return_normal: Whether to return the ray hit normals. Defaults to False`. + return_normal: Whether to return the ray hit normals. Defaults to False. return_face_id: Whether to return the ray hit face ids. Defaults to False. return_mesh_id: Whether to return the mesh id. Defaults to False. """ # get the thread id tid_mesh_id, tid_env, tid_ray = wp.tid() + if not env_mask[tid_env]: + return mesh_pose = wp.transform(mesh_positions[tid_env, tid_mesh_id], mesh_rotations[tid_env, tid_mesh_id]) mesh_pose_inv = wp.transform_inverse(mesh_pose) @@ -225,10 +306,11 @@ def raycast_dynamic_meshes_kernel( # if the ray hit, store the hit data if mesh_query_ray_t.result: wp.atomic_min(ray_distance, tid_env, tid_ray, mesh_query_ray_t.t) - # check if hit distance is less than the current hit distance, only then update the memory - # TODO, in theory we could use the output of atomic_min to avoid the non-thread safe next comparison - # however, warp atomic_min is returning the wrong values on gpu currently. - # FIXME https://github.com/NVIDIA/warp/issues/1058 + # TODO(warp#1058): Use the return value of atomic_min to avoid the non-thread-safe + # equality check below. Currently warp atomic_min returns wrong values on GPU, so we + # fall back to a racy read-back. When two meshes tie on distance, normals/face-ids/ + # mesh-ids may be written by different threads. The hit *position* is still correct + # because all tying threads compute the same world-space point. if mesh_query_ray_t.t == ray_distance[tid_env, tid_ray]: # convert back to world space and update the hit data hit_pos = start_pos + mesh_query_ray_t.t * direction @@ -302,379 +384,324 @@ def reshape_tiled_image( ) ## -# Wrench Composer +# Wrench Composer — Dual-Buffer Architecture ## -@wp.func -def cast_to_link_frame(position: wp.vec3f, link_position: wp.vec3f, is_global: bool) -> wp.vec3f: - """Casts a position to the link frame of the body. - - Args: - position: The position to cast. - link_position: The link frame position. - is_global: Whether the position is in the global frame. - - Returns: - The position in the link frame of the body. - """ - if is_global: - return position - link_position - else: - return position - - -@wp.func -def cast_force_to_link_frame(force: wp.vec3f, link_quat: wp.quatf, is_global: bool) -> wp.vec3f: - """Casts a force to the link frame of the body. - - Args: - force: The force to cast. - link_quat: The link frame quaternion. - is_global: Whether the force is applied in the global frame. - Returns: - The force in the link frame of the body. - """ - if is_global: - return wp.quat_rotate_inv(link_quat, force) - else: - return force - - -@wp.func -def cast_torque_to_link_frame(torque: wp.vec3f, link_quat: wp.quatf, is_global: bool) -> wp.vec3f: - """Casts a torque to the link frame of the body. - - Args: - torque: The torque to cast. - link_quat: The link frame quaternion. - is_global: Whether the torque is applied in the global frame. - - Returns: - The torque in the link frame of the body. - """ - if is_global: - return wp.quat_rotate_inv(link_quat, torque) - else: - return torque - - @wp.kernel -def add_forces_and_torques_at_position_index( +def set_forces_to_dual_buffers_index( env_ids: wp.array(dtype=wp.int32), body_ids: wp.array(dtype=wp.int32), forces: wp.array2d(dtype=wp.vec3f), torques: wp.array2d(dtype=wp.vec3f), positions: wp.array2d(dtype=wp.vec3f), - link_poses: wp.array2d(dtype=wp.transformf), + global_force_w: wp.array2d(dtype=wp.vec3f), + global_torque_w: wp.array2d(dtype=wp.vec3f), + global_force_at_com_w: wp.array2d(dtype=wp.vec3f), + local_force_b: wp.array2d(dtype=wp.vec3f), + local_torque_b: wp.array2d(dtype=wp.vec3f), is_global: bool, - composed_forces_b: wp.array2d(dtype=wp.vec3f), - composed_torques_b: wp.array2d(dtype=wp.vec3f), ): - """Add forces and torques to the composed wrench at user-provided positions using index selection. + """Set forces/torques into dual buffers using index selection (overwrites). - When is_global is False, the user-provided positions offset the force application relative to - the link frame. When is_global is True, positions are in the global frame. Results are - accumulated (added) into the composed buffers. + Dispatched with ``dim=(len(env_ids), len(body_ids))``. - .. note:: - Expects partial data from the user (indexed by env_ids/body_ids). + When ``is_global`` is True, forces/torques are written to the world-frame buffers. + Forces with ``positions`` go to ``global_force_w`` with torque ``cross(P, F)`` accumulated + into ``global_torque_w``; forces without positions go to ``global_force_at_com_w``. + When ``is_global`` is False, values go to ``local_force_b`` / ``local_torque_b``. - Args: - env_ids: Input array of environment indices. Shape is (num_selected_envs,). - body_ids: Input array of body indices. Shape is (num_selected_bodies,). - forces: Input array of forces to apply. Shape is (num_selected_envs, num_selected_bodies). - Can be None if not provided. - torques: Input array of torques to apply. Shape is (num_selected_envs, num_selected_bodies). - Can be None if not provided. - positions: Input array of position offsets for force application. - Shape is (num_selected_envs, num_selected_bodies). Can be None if not provided. - link_poses: Input array of link frame poses in world frame. - Shape is (num_envs, num_bodies). - is_global: Input flag indicating whether forces/torques/positions are in the global frame. - composed_forces_b: Output array where forces in the link frame are accumulated. - Shape is (num_envs, num_bodies). - composed_torques_b: Output array where torques in the link frame are accumulated. - Shape is (num_envs, num_bodies). + Any of ``forces``, ``torques``, or ``positions`` may be ``None`` (null array). """ - # get the thread id tid_env, tid_body = wp.tid() + ei = env_ids[tid_env] + bi = body_ids[tid_body] - # add the forces to the composed force, if the positions are provided, also adds a torque to the composed torque. - if forces: - # add the forces to the composed force - composed_forces_b[env_ids[tid_env], body_ids[tid_body]] += cast_force_to_link_frame( - forces[tid_env, tid_body], - wp.transform_get_rotation(link_poses[env_ids[tid_env], body_ids[tid_body]]), - is_global, - ) - # if there is a position offset, add a torque to the composed torque. - if positions: - composed_torques_b[env_ids[tid_env], body_ids[tid_body]] += wp.skew( - cast_to_link_frame( - positions[tid_env, tid_body], - wp.transform_get_translation(link_poses[env_ids[tid_env], body_ids[tid_body]]), - is_global, - ) - ) @ cast_force_to_link_frame( - forces[tid_env, tid_body], - wp.transform_get_rotation(link_poses[env_ids[tid_env], body_ids[tid_body]]), - is_global, - ) - if torques: - composed_torques_b[env_ids[tid_env], body_ids[tid_body]] += cast_torque_to_link_frame( - torques[tid_env, tid_body], - wp.transform_get_rotation(link_poses[env_ids[tid_env], body_ids[tid_body]]), - is_global, - ) + if is_global: + if torques: + global_torque_w[ei, bi] = torques[tid_env, tid_body] + if forces: + if positions: + global_force_w[ei, bi] = forces[tid_env, tid_body] + if torques: + global_torque_w[ei, bi] = global_torque_w[ei, bi] + wp.cross( + positions[tid_env, tid_body], forces[tid_env, tid_body] + ) + else: + global_torque_w[ei, bi] = wp.cross(positions[tid_env, tid_body], forces[tid_env, tid_body]) + else: + global_force_at_com_w[ei, bi] = forces[tid_env, tid_body] + else: + if torques: + local_torque_b[ei, bi] = torques[tid_env, tid_body] + if forces: + local_force_b[ei, bi] = forces[tid_env, tid_body] + if positions: + if torques: + local_torque_b[ei, bi] = local_torque_b[ei, bi] + wp.cross( + positions[tid_env, tid_body], forces[tid_env, tid_body] + ) + else: + local_torque_b[ei, bi] = wp.cross(positions[tid_env, tid_body], forces[tid_env, tid_body]) @wp.kernel -def set_forces_and_torques_at_position_index( +def add_forces_to_dual_buffers_index( env_ids: wp.array(dtype=wp.int32), body_ids: wp.array(dtype=wp.int32), forces: wp.array2d(dtype=wp.vec3f), torques: wp.array2d(dtype=wp.vec3f), positions: wp.array2d(dtype=wp.vec3f), - link_poses: wp.array2d(dtype=wp.transformf), + global_force_w: wp.array2d(dtype=wp.vec3f), + global_torque_w: wp.array2d(dtype=wp.vec3f), + global_force_at_com_w: wp.array2d(dtype=wp.vec3f), + local_force_b: wp.array2d(dtype=wp.vec3f), + local_torque_b: wp.array2d(dtype=wp.vec3f), is_global: bool, - composed_forces_b: wp.array2d(dtype=wp.vec3f), - composed_torques_b: wp.array2d(dtype=wp.vec3f), ): - """Set forces and torques to the composed wrench at user-provided positions using index selection. - - When is_global is False, the user-provided positions offset the force application relative to - the link frame. When is_global is True, positions are in the global frame. Results are - overwritten (set) in the composed buffers. + """Add forces/torques into dual buffers using index selection (accumulates). - .. note:: - Expects partial data from the user (indexed by env_ids/body_ids). - - Args: - env_ids: Input array of environment indices. Shape is (num_selected_envs,). - body_ids: Input array of body indices. Shape is (num_selected_bodies,). - forces: Input array of forces to apply. Shape is (num_selected_envs, num_selected_bodies). - Can be None if not provided. - torques: Input array of torques to apply. Shape is (num_selected_envs, num_selected_bodies). - Can be None if not provided. - positions: Input array of position offsets for force application. - Shape is (num_selected_envs, num_selected_bodies). Can be None if not provided. - link_poses: Input array of link frame poses in world frame. - Shape is (num_envs, num_bodies). - is_global: Input flag indicating whether forces/torques/positions are in the global frame. - composed_forces_b: Output array where forces in the link frame are written. - Shape is (num_envs, num_bodies). - composed_torques_b: Output array where torques in the link frame are written. - Shape is (num_envs, num_bodies). + Same routing logic as :func:`set_forces_to_dual_buffers_index` but uses ``+=`` instead of ``=``. + Dispatched with ``dim=(len(env_ids), len(body_ids))``. """ - # get the thread id tid_env, tid_body = wp.tid() + ei = env_ids[tid_env] + bi = body_ids[tid_body] - # set the torques to the composed torque - if torques: - composed_torques_b[env_ids[tid_env], body_ids[tid_body]] = cast_torque_to_link_frame( - torques[tid_env, tid_body], - wp.transform_get_rotation(link_poses[env_ids[tid_env], body_ids[tid_body]]), - is_global, - ) - # set the forces to the composed force, if the positions are provided, adds a torque to the composed torque - # from the force at that position. - if forces: - # set the forces to the composed force - composed_forces_b[env_ids[tid_env], body_ids[tid_body]] = cast_force_to_link_frame( - forces[tid_env, tid_body], - wp.transform_get_rotation(link_poses[env_ids[tid_env], body_ids[tid_body]]), - is_global, - ) - # if there is a position offset, set the torque from the force at that position. - if positions: - composed_torques_b[env_ids[tid_env], body_ids[tid_body]] = wp.skew( - cast_to_link_frame( - positions[tid_env, tid_body], - wp.transform_get_translation(link_poses[env_ids[tid_env], body_ids[tid_body]]), - is_global, + if is_global: + if forces: + if positions: + global_force_w[ei, bi] = global_force_w[ei, bi] + forces[tid_env, tid_body] + global_torque_w[ei, bi] = global_torque_w[ei, bi] + wp.cross( + positions[tid_env, tid_body], forces[tid_env, tid_body] ) - ) @ cast_force_to_link_frame( - forces[tid_env, tid_body], - wp.transform_get_rotation(link_poses[env_ids[tid_env], body_ids[tid_body]]), - is_global, - ) + else: + global_force_at_com_w[ei, bi] = global_force_at_com_w[ei, bi] + forces[tid_env, tid_body] + if torques: + global_torque_w[ei, bi] = global_torque_w[ei, bi] + torques[tid_env, tid_body] + else: + if forces: + local_force_b[ei, bi] = local_force_b[ei, bi] + forces[tid_env, tid_body] + if positions: + local_torque_b[ei, bi] = local_torque_b[ei, bi] + wp.cross( + positions[tid_env, tid_body], forces[tid_env, tid_body] + ) + if torques: + local_torque_b[ei, bi] = local_torque_b[ei, bi] + torques[tid_env, tid_body] @wp.kernel -def add_forces_and_torques_at_position_mask( +def set_forces_to_dual_buffers_mask( env_mask: wp.array(dtype=wp.bool), body_mask: wp.array(dtype=wp.bool), forces: wp.array2d(dtype=wp.vec3f), torques: wp.array2d(dtype=wp.vec3f), positions: wp.array2d(dtype=wp.vec3f), - link_poses: wp.array2d(dtype=wp.transformf), + global_force_w: wp.array2d(dtype=wp.vec3f), + global_torque_w: wp.array2d(dtype=wp.vec3f), + global_force_at_com_w: wp.array2d(dtype=wp.vec3f), + local_force_b: wp.array2d(dtype=wp.vec3f), + local_torque_b: wp.array2d(dtype=wp.vec3f), is_global: bool, - composed_forces_b: wp.array2d(dtype=wp.vec3f), - composed_torques_b: wp.array2d(dtype=wp.vec3f), ): - """Add forces and torques to the composed wrench at user-provided positions using mask selection. - - When is_global is False, the user-provided positions offset the force application relative to - the link frame. When is_global is True, positions are in the global frame. Results are - accumulated (added) into the composed buffers. Only entries where both env_mask and body_mask - are True are processed. - - .. note:: - Expects full data from the user (num_envs x num_bodies). + """Set forces/torques into dual buffers using mask selection (overwrites). - Args: - env_mask: Input boolean mask for environments. Shape is (num_envs,). - body_mask: Input boolean mask for bodies. Shape is (num_bodies,). - forces: Input array of forces to apply. Shape is (num_envs, num_bodies). - Can be None if not provided. - torques: Input array of torques to apply. Shape is (num_envs, num_bodies). - Can be None if not provided. - positions: Input array of position offsets for force application. - Shape is (num_envs, num_bodies). Can be None if not provided. - link_poses: Input array of link frame poses in world frame. - Shape is (num_envs, num_bodies). - is_global: Input flag indicating whether forces/torques/positions are in the global frame. - composed_forces_b: Output array where forces in the link frame are accumulated. - Shape is (num_envs, num_bodies). - composed_torques_b: Output array where torques in the link frame are accumulated. - Shape is (num_envs, num_bodies). + Same routing logic as :func:`set_forces_to_dual_buffers_index` but threads are gated by + ``env_mask[tid_env] and body_mask[tid_body]``, and indices are direct (no indirection array). + Dispatched with ``dim=(num_envs, num_bodies)``. """ - # get the thread id tid_env, tid_body = wp.tid() if env_mask[tid_env] and body_mask[tid_body]: - # add the forces to the composed force, if the positions are provided, also adds a torque to the composed - # torque. - if forces: - # add the forces to the composed force - composed_forces_b[tid_env, tid_body] += cast_force_to_link_frame( - forces[tid_env, tid_body], wp.transform_get_rotation(link_poses[tid_env, tid_body]), is_global - ) - # if there is a position offset, add a torque to the composed torque. - if positions: - composed_torques_b[tid_env, tid_body] += wp.skew( - cast_to_link_frame( - positions[tid_env, tid_body], - wp.transform_get_translation(link_poses[tid_env, tid_body]), - is_global, - ) - ) @ cast_force_to_link_frame( - forces[tid_env, tid_body], wp.transform_get_rotation(link_poses[tid_env, tid_body]), is_global - ) - if torques: - composed_torques_b[tid_env, tid_body] += cast_torque_to_link_frame( - torques[tid_env, tid_body], wp.transform_get_rotation(link_poses[tid_env, tid_body]), is_global - ) + if is_global: + if torques: + global_torque_w[tid_env, tid_body] = torques[tid_env, tid_body] + if forces: + if positions: + global_force_w[tid_env, tid_body] = forces[tid_env, tid_body] + if torques: + global_torque_w[tid_env, tid_body] = global_torque_w[tid_env, tid_body] + wp.cross( + positions[tid_env, tid_body], forces[tid_env, tid_body] + ) + else: + global_torque_w[tid_env, tid_body] = wp.cross( + positions[tid_env, tid_body], forces[tid_env, tid_body] + ) + else: + global_force_at_com_w[tid_env, tid_body] = forces[tid_env, tid_body] + else: + if torques: + local_torque_b[tid_env, tid_body] = torques[tid_env, tid_body] + if forces: + local_force_b[tid_env, tid_body] = forces[tid_env, tid_body] + if positions: + if torques: + local_torque_b[tid_env, tid_body] = local_torque_b[tid_env, tid_body] + wp.cross( + positions[tid_env, tid_body], forces[tid_env, tid_body] + ) + else: + local_torque_b[tid_env, tid_body] = wp.cross( + positions[tid_env, tid_body], forces[tid_env, tid_body] + ) @wp.kernel -def set_forces_and_torques_at_position_mask( +def add_forces_to_dual_buffers_mask( env_mask: wp.array(dtype=wp.bool), body_mask: wp.array(dtype=wp.bool), forces: wp.array2d(dtype=wp.vec3f), torques: wp.array2d(dtype=wp.vec3f), positions: wp.array2d(dtype=wp.vec3f), - link_poses: wp.array2d(dtype=wp.transformf), + global_force_w: wp.array2d(dtype=wp.vec3f), + global_torque_w: wp.array2d(dtype=wp.vec3f), + global_force_at_com_w: wp.array2d(dtype=wp.vec3f), + local_force_b: wp.array2d(dtype=wp.vec3f), + local_torque_b: wp.array2d(dtype=wp.vec3f), is_global: bool, - composed_forces_b: wp.array2d(dtype=wp.vec3f), - composed_torques_b: wp.array2d(dtype=wp.vec3f), ): - """Set forces and torques to the composed wrench at user-provided positions using mask selection. + """Add forces/torques into dual buffers using mask selection (accumulates). - When is_global is False, the user-provided positions offset the force application relative to - the link frame. When is_global is True, positions are in the global frame. Results are - overwritten (set) in the composed buffers. Only entries where both env_mask and body_mask - are True are processed. - - .. note:: - Expects full data from the user (num_envs x num_bodies). - - Args: - env_mask: Input boolean mask for environments. Shape is (num_envs,). - body_mask: Input boolean mask for bodies. Shape is (num_bodies,). - forces: Input array of forces to apply. Shape is (num_envs, num_bodies). - Can be None if not provided. - torques: Input array of torques to apply. Shape is (num_envs, num_bodies). - Can be None if not provided. - positions: Input array of position offsets for force application. - Shape is (num_envs, num_bodies). Can be None if not provided. - link_poses: Input array of link frame poses in world frame. - Shape is (num_envs, num_bodies). - is_global: Input flag indicating whether forces/torques/positions are in the global frame. - composed_forces_b: Output array where forces in the link frame are written. - Shape is (num_envs, num_bodies). - composed_torques_b: Output array where torques in the link frame are written. - Shape is (num_envs, num_bodies). + Same routing logic as :func:`add_forces_to_dual_buffers_index` but threads are gated by + ``env_mask[tid_env] and body_mask[tid_body]``. + Dispatched with ``dim=(num_envs, num_bodies)``. """ - # get the thread id tid_env, tid_body = wp.tid() - # set the torques to the composed torque if env_mask[tid_env] and body_mask[tid_body]: - if torques: - composed_torques_b[tid_env, tid_body] = cast_torque_to_link_frame( - torques[tid_env, tid_body], wp.transform_get_rotation(link_poses[tid_env, tid_body]), is_global - ) - # set the forces to the composed force, if the positions are provided, adds a torque to the composed torque - # from the force at that position. - if forces: - # set the forces to the composed force - composed_forces_b[tid_env, tid_body] = cast_force_to_link_frame( - forces[tid_env, tid_body], wp.transform_get_rotation(link_poses[tid_env, tid_body]), is_global - ) - # if there is a position offset, set the torque from the force at that position. - if positions: - composed_torques_b[tid_env, tid_body] = wp.skew( - cast_to_link_frame( - positions[tid_env, tid_body], - wp.transform_get_translation(link_poses[tid_env, tid_body]), - is_global, + if is_global: + if forces: + if positions: + global_force_w[tid_env, tid_body] = global_force_w[tid_env, tid_body] + forces[tid_env, tid_body] + global_torque_w[tid_env, tid_body] = global_torque_w[tid_env, tid_body] + wp.cross( + positions[tid_env, tid_body], forces[tid_env, tid_body] ) - ) @ cast_force_to_link_frame( - forces[tid_env, tid_body], wp.transform_get_rotation(link_poses[tid_env, tid_body]), is_global - ) + else: + global_force_at_com_w[tid_env, tid_body] = ( + global_force_at_com_w[tid_env, tid_body] + forces[tid_env, tid_body] + ) + if torques: + global_torque_w[tid_env, tid_body] = global_torque_w[tid_env, tid_body] + torques[tid_env, tid_body] + else: + if forces: + local_force_b[tid_env, tid_body] = local_force_b[tid_env, tid_body] + forces[tid_env, tid_body] + if positions: + local_torque_b[tid_env, tid_body] = local_torque_b[tid_env, tid_body] + wp.cross( + positions[tid_env, tid_body], forces[tid_env, tid_body] + ) + if torques: + local_torque_b[tid_env, tid_body] = local_torque_b[tid_env, tid_body] + torques[tid_env, tid_body] @wp.kernel -def reset_wrench_composer_index( - env_ids: wp.array(dtype=wp.int32), - composed_forces_b: wp.array2d(dtype=wp.vec3f), - composed_torques_b: wp.array2d(dtype=wp.vec3f), +def add_raw_wrench_buffers( + src_gf: wp.array2d(dtype=wp.vec3f), + src_gt: wp.array2d(dtype=wp.vec3f), + src_gfc: wp.array2d(dtype=wp.vec3f), + src_lf: wp.array2d(dtype=wp.vec3f), + src_lt: wp.array2d(dtype=wp.vec3f), + dst_gf: wp.array2d(dtype=wp.vec3f), + dst_gt: wp.array2d(dtype=wp.vec3f), + dst_gfc: wp.array2d(dtype=wp.vec3f), + dst_lf: wp.array2d(dtype=wp.vec3f), + dst_lt: wp.array2d(dtype=wp.vec3f), ): - """Reset the composed force and torque to zero at the specified environment indices. + """Element-wise add all five source wrench buffers into destination buffers. - Args: - env_ids: Input array of environment indices to reset. Shape is (num_selected_envs,). - composed_forces_b: Output array where forces are zeroed. Shape is (num_envs, num_bodies). - composed_torques_b: Output array where torques are zeroed. Shape is (num_envs, num_bodies). + Dispatched with ``dim=(num_envs, num_bodies)``. Each ``src_*`` / ``dst_*`` pair corresponds + to one of the five input buffers (global_force_w, global_torque_w, global_force_at_com_w, + local_force_b, local_torque_b). """ + tid_env, tid_body = wp.tid() + dst_gf[tid_env, tid_body] = dst_gf[tid_env, tid_body] + src_gf[tid_env, tid_body] + dst_gt[tid_env, tid_body] = dst_gt[tid_env, tid_body] + src_gt[tid_env, tid_body] + dst_gfc[tid_env, tid_body] = dst_gfc[tid_env, tid_body] + src_gfc[tid_env, tid_body] + dst_lf[tid_env, tid_body] = dst_lf[tid_env, tid_body] + src_lf[tid_env, tid_body] + dst_lt[tid_env, tid_body] = dst_lt[tid_env, tid_body] + src_lt[tid_env, tid_body] - # get the thread id + +@wp.kernel +def compose_wrench_to_body_frame( + global_force_w: wp.array2d(dtype=wp.vec3f), + global_torque_w: wp.array2d(dtype=wp.vec3f), + global_force_at_com_w: wp.array2d(dtype=wp.vec3f), + local_force_b: wp.array2d(dtype=wp.vec3f), + local_torque_b: wp.array2d(dtype=wp.vec3f), + com_pos_w: wp.array2d(dtype=wp.vec3f), + link_quat_w: wp.array2d(dtype=wp.quatf), + out_force_b: wp.array2d(dtype=wp.vec3f), + out_torque_b: wp.array2d(dtype=wp.vec3f), +): + """Compose global and local wrench buffers into a single body-frame output. + + Global torques store the moment of positional forces about the world origin: ``cross(P, F)``. + This kernel corrects to be about the body's CoM via ``cross(P, F) - cross(com_pos_w, F) = + cross(P - com_pos_w, F)``, then rotates both force and torque into the body frame using + ``quat_rotate_inv(link_quat_w, ...)``, and adds local-frame values. + + Dispatched with ``dim=(num_envs, num_bodies)``. + """ tid_env, tid_body = wp.tid() + total_force_w = global_force_w[tid_env, tid_body] + global_force_at_com_w[tid_env, tid_body] + corrected_torque_w = global_torque_w[tid_env, tid_body] - wp.cross( + com_pos_w[tid_env, tid_body], global_force_w[tid_env, tid_body] + ) + out_force_b[tid_env, tid_body] = ( + wp.quat_rotate_inv(link_quat_w[tid_env, tid_body], total_force_w) + local_force_b[tid_env, tid_body] + ) + out_torque_b[tid_env, tid_body] = ( + wp.quat_rotate_inv(link_quat_w[tid_env, tid_body], corrected_torque_w) + local_torque_b[tid_env, tid_body] + ) + + +@wp.kernel +def reset_wrench_composer_index( + env_ids: wp.array(dtype=wp.int32), + global_force_w: wp.array2d(dtype=wp.vec3f), + global_torque_w: wp.array2d(dtype=wp.vec3f), + global_force_at_com_w: wp.array2d(dtype=wp.vec3f), + local_force_b: wp.array2d(dtype=wp.vec3f), + local_torque_b: wp.array2d(dtype=wp.vec3f), + out_force_b: wp.array2d(dtype=wp.vec3f), + out_torque_b: wp.array2d(dtype=wp.vec3f), +): + """Zero all 7 wrench composer buffers at the specified environment indices. - # reset the composed force and torque - composed_forces_b[env_ids[tid_env], tid_body] = wp.vec3f(0.0) - composed_torques_b[env_ids[tid_env], tid_body] = wp.vec3f(0.0) + Dispatched with ``dim=(len(env_ids), num_bodies)``. + """ + tid_env, tid_body = wp.tid() + ei = env_ids[tid_env] + z = wp.vec3f(0.0) + global_force_w[ei, tid_body] = z + global_torque_w[ei, tid_body] = z + global_force_at_com_w[ei, tid_body] = z + local_force_b[ei, tid_body] = z + local_torque_b[ei, tid_body] = z + out_force_b[ei, tid_body] = z + out_torque_b[ei, tid_body] = z @wp.kernel def reset_wrench_composer_mask( env_mask: wp.array(dtype=wp.bool), - composed_forces_b: wp.array2d(dtype=wp.vec3f), - composed_torques_b: wp.array2d(dtype=wp.vec3f), + global_force_w: wp.array2d(dtype=wp.vec3f), + global_torque_w: wp.array2d(dtype=wp.vec3f), + global_force_at_com_w: wp.array2d(dtype=wp.vec3f), + local_force_b: wp.array2d(dtype=wp.vec3f), + local_torque_b: wp.array2d(dtype=wp.vec3f), + out_force_b: wp.array2d(dtype=wp.vec3f), + out_torque_b: wp.array2d(dtype=wp.vec3f), ): - """Reset the composed force and torque to zero for environments matching the mask. + """Zero all 7 wrench composer buffers for environments matching the mask. - Args: - env_mask: Input boolean mask for environments. Shape is (num_envs,). - composed_forces_b: Output array where forces are zeroed. Shape is (num_envs, num_bodies). - composed_torques_b: Output array where torques are zeroed. Shape is (num_envs, num_bodies). + Dispatched with ``dim=(num_envs, num_bodies)``. """ - # get the thread id tid_env, tid_body = wp.tid() - - # reset the composed force and torque if env_mask[tid_env]: - composed_forces_b[tid_env, tid_body] = wp.vec3f(0.0) - composed_torques_b[tid_env, tid_body] = wp.vec3f(0.0) + z = wp.vec3f(0.0) + global_force_w[tid_env, tid_body] = z + global_torque_w[tid_env, tid_body] = z + global_force_at_com_w[tid_env, tid_body] = z + local_force_b[tid_env, tid_body] = z + local_torque_b[tid_env, tid_body] = z + out_force_b[tid_env, tid_body] = z + out_torque_b[tid_env, tid_body] = z diff --git a/source/isaaclab/isaaclab/utils/warp/ops.py b/source/isaaclab/isaaclab/utils/warp/ops.py index 313a7fd43afb..a3ee273b627c 100644 --- a/source/isaaclab/isaaclab/utils/warp/ops.py +++ b/source/isaaclab/isaaclab/utils/warp/ops.py @@ -19,6 +19,10 @@ from . import kernels +# Cache of all-True env masks keyed by (n_envs, device) to avoid per-call allocations in +# raycast_dynamic_meshes. Populated lazily on first call with a given (n_envs, device) pair. +_all_env_mask_cache: dict[tuple[int, str], wp.array] = {} + def raycast_mesh( ray_starts: torch.Tensor, @@ -335,11 +339,19 @@ def raycast_dynamic_meshes( mesh_orientations_w = mesh_orientations_w.to(dtype=torch.float32, device=torch_device).contiguous() mesh_quat_wp_w = wp.from_torch(mesh_orientations_w, dtype=wp.quat) + # All environments active when called through this public API. + # Cache the mask by (n_envs, device) to avoid a per-call allocation. + cache_key = (n_envs, str(torch_device)) + if cache_key not in _all_env_mask_cache: + _all_env_mask_cache[cache_key] = wp.from_torch(torch.ones(n_envs, dtype=torch.bool, device=torch_device)) + all_env_mask = _all_env_mask_cache[cache_key] + # launch the warp kernel wp.launch( kernel=kernels.raycast_dynamic_meshes_kernel, dim=[n_meshes, n_envs, n_rays_per_env], inputs=[ + all_env_mask, mesh_ids_wp, ray_starts_wp, ray_directions_wp, diff --git a/source/isaaclab/isaaclab/utils/wrench_composer.py b/source/isaaclab/isaaclab/utils/wrench_composer.py index 5ad966a6e4e9..e348697306d5 100644 --- a/source/isaaclab/isaaclab/utils/wrench_composer.py +++ b/source/isaaclab/isaaclab/utils/wrench_composer.py @@ -6,6 +6,7 @@ from __future__ import annotations import warnings +from collections.abc import Sequence from typing import TYPE_CHECKING import numpy as np @@ -13,12 +14,14 @@ import warp as wp from isaaclab.utils.warp.kernels import ( - add_forces_and_torques_at_position_index, - add_forces_and_torques_at_position_mask, + add_forces_to_dual_buffers_index, + add_forces_to_dual_buffers_mask, + add_raw_wrench_buffers, + compose_wrench_to_body_frame, reset_wrench_composer_index, reset_wrench_composer_mask, - set_forces_and_torques_at_position_index, - set_forces_and_torques_at_position_mask, + set_forces_to_dual_buffers_index, + set_forces_to_dual_buffers_mask, ) if TYPE_CHECKING: @@ -27,16 +30,33 @@ class WrenchComposer: def __init__(self, asset: BaseArticulation | BaseRigidObject | BaseRigidObjectCollection) -> None: - """Wrench composer. + """Wrench composer with dual-buffer architecture. - This class is used to compose forces and torques at the body's link frame. - It can compose global wrenches and local wrenches. The result is always in the link frame of the body. + This class composes forces and torques applied to rigid bodies. Forces and torques can be + specified in either the global (world) frame or the local (body) frame. Internally, they are + stored in separate global and local input buffers. When the final composed wrench is needed, + the global contributions are rotated into the body frame and combined with the local + contributions to produce the output force and torque expressed in the body frame. + + The dual-buffer architecture uses five input buffers: + + - ``global_force_w``: Global forces [N] (world frame). + - ``global_torque_w``: Global torques [N·m] (world frame), including moment contributions + from positional forces (``cross(P, F)``). + - ``global_force_at_com_w``: Global forces [N] applied at the body's CoM (world frame, no positional torque). + - ``local_force_b``: Local forces [N] (body frame). + - ``local_torque_b``: Local torques [N·m] (body frame). + + And two output buffers: + + - ``out_force_b``: Composed force [N] in body frame. + - ``out_torque_b``: Composed torque [N·m] in body frame. Args: - asset: Asset to use. Defaults to None. + asset: Asset to use. """ self.num_envs = asset.num_instances - # Avoid isinstance to prevent circular import issues, use attribute presence instead. + # Avoid isinstance to prevent circular import issues; check by attribute presence instead. if hasattr(asset, "num_bodies"): self.num_bodies = asset.num_bodies else: @@ -44,16 +64,28 @@ def __init__(self, asset: BaseArticulation | BaseRigidObject | BaseRigidObjectCo self.device = asset.device self._asset = asset self._active = False - - # Avoid isinstance here due to potential circular import issues; check by attribute presence instead. - if hasattr(self._asset.data, "body_link_pose_w"): - self._get_link_pose_fn = lambda a=self._asset: a.data.body_link_pose_w + self._dirty = False + if hasattr(self._asset.data, "body_com_pos_w"): + self._get_com_pos_fn = lambda a=self._asset: a.data.body_com_pos_w else: raise ValueError(f"Unsupported asset type: {self._asset.__class__.__name__}") + if hasattr(self._asset.data, "body_link_quat_w"): + self._get_link_quat_fn = lambda a=self._asset: a.data.body_link_quat_w + else: + raise ValueError(f"Unsupported asset type: {self._asset.__class__.__name__}") + + # -- Input buffers (5 total) -- + self._global_force_w = wp.zeros((self.num_envs, self.num_bodies), dtype=wp.vec3f, device=self.device) + self._global_torque_w = wp.zeros((self.num_envs, self.num_bodies), dtype=wp.vec3f, device=self.device) + self._global_force_at_com_w = wp.zeros((self.num_envs, self.num_bodies), dtype=wp.vec3f, device=self.device) + self._local_force_b = wp.zeros((self.num_envs, self.num_bodies), dtype=wp.vec3f, device=self.device) + self._local_torque_b = wp.zeros((self.num_envs, self.num_bodies), dtype=wp.vec3f, device=self.device) + + # -- Output buffers (2 total) -- + self._out_force_b = wp.zeros((self.num_envs, self.num_bodies), dtype=wp.vec3f, device=self.device) + self._out_torque_b = wp.zeros((self.num_envs, self.num_bodies), dtype=wp.vec3f, device=self.device) - # Create buffers - self._composed_force_b = wp.zeros((self.num_envs, self.num_bodies), dtype=wp.vec3f, device=self.device) - self._composed_torque_b = wp.zeros((self.num_envs, self.num_bodies), dtype=wp.vec3f, device=self.device) + # -- Index / mask helper arrays -- self._ALL_ENV_INDICES = wp.array(np.arange(self.num_envs, dtype=np.int32), dtype=wp.int32, device=self.device) self._ALL_BODY_INDICES = wp.array( np.arange(self.num_bodies, dtype=np.int32), dtype=wp.int32, device=self.device @@ -61,44 +93,125 @@ def __init__(self, asset: BaseArticulation | BaseRigidObject | BaseRigidObjectCo self._ALL_ENV_MASK = wp.ones((self.num_envs), dtype=wp.bool, device=self.device) self._ALL_BODY_MASK = wp.ones((self.num_bodies), dtype=wp.bool, device=self.device) - # Temporary buffers for the masks, positions, and forces/torques (reused to avoid allocations) - self._temp_env_mask_wp = wp.zeros((self.num_envs,), dtype=wp.bool, device=self.device) - self._temp_body_mask_wp = wp.zeros((self.num_bodies,), dtype=wp.bool, device=self.device) - self._temp_positions_wp = wp.zeros((self.num_envs, self.num_bodies), dtype=wp.vec3f, device=self.device) - self._temp_forces_wp = wp.zeros((self.num_envs, self.num_bodies), dtype=wp.vec3f, device=self.device) - self._temp_torques_wp = wp.zeros((self.num_envs, self.num_bodies), dtype=wp.vec3f, device=self.device) - - # Flag to check if the link poses have been updated. - self._link_poses_updated = False + # ------------------------------------------------------------------ + # Properties + # ------------------------------------------------------------------ @property def active(self) -> bool: - """Whether the wrench composer is active.""" + """Whether the wrench composer is active (has pending forces/torques). + + Set to ``True`` when any ``add_*`` or ``set_*`` method writes data. Cleared only by a + full :meth:`reset` call (no arguments). Partial resets (with ``env_ids`` or ``env_mask``) + do **not** clear this flag because checking whether all environments are zero would + require scanning the buffers, defeating the purpose of a cheap guard. + + This means the flag may remain ``True`` even if all buffers are zero after partial resets. + This is by design: the cost of an unnecessary compose + apply on zero data is negligible + compared to scanning the buffers every frame. + """ return self._active @property - def composed_force(self) -> wp.array: - """Composed force at the body's link frame. + def global_force_w(self) -> wp.array: + """Global force buffer [N] (world frame), dtype ``wp.vec3f``. Shape: ``(num_envs, num_bodies)``. - .. note:: If some of the forces are applied in the global frame, the composed force will be in the link frame - of the body. + .. note:: + This returns the underlying buffer reference for read-only inspection. Writing to it + directly bypasses the dirty flag and may produce stale output buffers. Use the + ``add_*`` or ``set_*`` methods to modify forces. + """ + return self._global_force_w - Returns: - wp.array: Composed force at the body's link frame. (num_envs, num_bodies, 3) + @property + def global_torque_w(self) -> wp.array: + """Global torque buffer [N·m] (world frame), dtype ``wp.vec3f``. Shape: ``(num_envs, num_bodies)``. + + Stores user-supplied torques plus moment contributions from positional forces (``cross(P, F)``). + + .. note:: + Read-only reference. See :attr:`global_force_w` for caveats on direct writes. """ - return self._composed_force_b + return self._global_torque_w @property - def composed_torque(self) -> wp.array: - """Composed torque at the body's link frame. + def global_force_at_com_w(self) -> wp.array: + """Global force at body's CoM buffer [N] (world frame, no positional torque). - .. note:: If some of the torques are applied in the global frame, the composed torque will be in the link frame - of the body. + dtype ``wp.vec3f``. Shape: ``(num_envs, num_bodies)``. - Returns: - wp.array: Composed torque at the body's link frame. (num_envs, num_bodies, 3) + .. note:: + Read-only reference. See :attr:`global_force_w` for caveats on direct writes. """ - return self._composed_torque_b + return self._global_force_at_com_w + + @property + def local_force_b(self) -> wp.array: + """Local force buffer [N] (body frame), dtype ``wp.vec3f``. Shape: ``(num_envs, num_bodies)``. + + .. note:: + Read-only reference. See :attr:`global_force_w` for caveats on direct writes. + """ + return self._local_force_b + + @property + def local_torque_b(self) -> wp.array: + """Local torque buffer [N·m] (body frame), dtype ``wp.vec3f``. Shape: ``(num_envs, num_bodies)``. + + .. note:: + Read-only reference. See :attr:`global_force_w` for caveats on direct writes. + """ + return self._local_torque_b + + @property + def out_force_b(self) -> wp.array: + """Composed output force [N] in the body frame, dtype ``wp.vec3f``. Shape: ``(num_envs, num_bodies)``. + + Triggers composition from input buffers if dirty. + """ + self._ensure_composed() + return self._out_force_b + + @property + def out_torque_b(self) -> wp.array: + """Composed output torque [N·m] in the body frame, dtype ``wp.vec3f``. Shape: ``(num_envs, num_bodies)``. + + Triggers composition from input buffers if dirty. + """ + self._ensure_composed() + return self._out_torque_b + + @property + def composed_force(self) -> wp.array: + """Composed force at the body frame, dtype ``wp.vec3f``. Shape: ``(num_envs, num_bodies)``. + + .. deprecated:: 4.5.33 + Use :attr:`out_force_b` instead. + """ + warnings.warn( + "The property 'composed_force' is deprecated. Use 'out_force_b' instead.", + DeprecationWarning, + stacklevel=2, + ) + return self.out_force_b + + @property + def composed_torque(self) -> wp.array: + """Composed torque at the body frame, dtype ``wp.vec3f``. Shape: ``(num_envs, num_bodies)``. + + .. deprecated:: 4.5.33 + Use :attr:`out_torque_b` instead. + """ + warnings.warn( + "The property 'composed_torque' is deprecated. Use 'out_torque_b' instead.", + DeprecationWarning, + stacklevel=2, + ) + return self.out_torque_b + + # ------------------------------------------------------------------ + # Public methods + # ------------------------------------------------------------------ def add_forces_and_torques_index( self, @@ -109,38 +222,26 @@ def add_forces_and_torques_index( env_ids: torch.Tensor | None = None, is_global: bool = False, ): - """Add forces and torques to the composed force and torque. - - Composed force and torque are the sum of all the forces and torques applied to the body. - It can compose global wrenches and local wrenches. The result is always in the link frame of the body. + """Add forces and torques into the input buffers using index-based selection. - The user can provide any combination of forces, torques, and positions. - - .. note:: Users may want to call `reset` function after every simulation step to ensure no force is carried - over to the next step. However, this may not necessary if the user calls `set_forces_and_torques` function - instead of `add_forces_and_torques`. + Accumulates onto whatever is already in the buffers. The result is always composed into the + body frame when the output properties are accessed. Args: - forces: Forces. (len(env_ids), len(body_ids), 3). Defaults to None. - torques: Torques. (len(env_ids), len(body_ids), 3). Defaults to None. - positions: Positions. (len(env_ids), len(body_ids), 3). Defaults to None. - body_ids: Body ids. Defaults to None (all bodies). - env_ids: Environment ids. Defaults to None (all environments). - is_global: Whether the forces and torques are applied in the global frame. Defaults to False. - - Raises: - ValueError: If the type of the input is not supported. - ValueError: If the input is a slice and it is not None. + forces: Forces [N]. Shape: (len(env_ids), len(body_ids), 3). Defaults to None. + torques: Torques [N·m]. Shape: (len(env_ids), len(body_ids), 3). Defaults to None. + positions: The positions [m] at which forces act. If `is_global` is True, these are global + positions expressed in the world frame. If `is_global` is False, these are offsets from the + body's CoM expressed in the body frame. If None, forces are assumed to act at the body's + CoM, independent of the `is_global` flag. + Shape: (len(env_ids), len(body_ids), 3). Defaults to None. + body_ids: Body indices. Defaults to None (all bodies). + env_ids: Environment indices. Defaults to None (all environments). + is_global: Whether the forces and torques are expressed in the global world frame or the local body frame. + Defaults to False. """ - # Resolve all indices - if (env_ids is None) or (env_ids == slice(None)): - env_ids = self._ALL_ENV_INDICES - if isinstance(env_ids, list): - env_ids = wp.array(env_ids, dtype=wp.int32, device=self.device) - if (body_ids is None) or (body_ids == slice(None)): - body_ids = self._ALL_BODY_INDICES - if isinstance(body_ids, list): - body_ids = wp.array(body_ids, dtype=wp.int32, device=self.device) + env_ids = self._resolve_env_ids(env_ids) + body_ids = self._resolve_body_ids(body_ids) if forces is None and torques is None: warnings.warn( "No forces or torques provided. No force will be added.", @@ -148,16 +249,12 @@ def add_forces_and_torques_index( stacklevel=2, ) return - # Get the link poses - if not self._link_poses_updated: - self._link_poses = self._get_link_pose_fn() - self._link_poses_updated = True - # Set the active flag to true self._active = True + self._dirty = True wp.launch( - add_forces_and_torques_at_position_index, + add_forces_to_dual_buffers_index, dim=(env_ids.shape[0], body_ids.shape[0]), inputs=[ env_ids, @@ -165,13 +262,13 @@ def add_forces_and_torques_index( forces, torques, positions, - self._link_poses, + self._global_force_w, + self._global_torque_w, + self._global_force_at_com_w, + self._local_force_b, + self._local_torque_b, is_global, ], - outputs=[ - self._composed_force_b, - self._composed_torque_b, - ], device=self.device, ) @@ -184,51 +281,43 @@ def set_forces_and_torques_index( env_ids: wp.array | torch.Tensor | None = None, is_global: bool = False, ): - """Set forces and torques to the composed force and torque. - - Composed force and torque are the sum of all the forces and torques applied to the body. - It can compose global wrenches and local wrenches. The result is always in the link frame of the body. + """Set forces and torques into the input buffers using index-based selection. - The user can provide any combination of forces, torques, and positions. + Resets the specified environments first, then writes the new values. This replaces any + previously accumulated forces/torques for the targeted environments while leaving other + environments untouched. Args: - forces: Forces. (num_envs, num_bodies, 3). Defaults to None. - torques: Torques. (num_envs, num_bodies, 3). Defaults to None. - positions: Positions. (num_envs, num_bodies, 3). Defaults to None. - body_ids: Body ids. (num_envs, num_bodies). Defaults to None (all bodies). - env_ids: Environment ids. (num_envs). Defaults to None (all environments). - is_global: Whether the forces and torques are applied in the global frame. Defaults to False. - - Raises: - ValueError: If the type of the input is not supported. - ValueError: If the input is a slice and it is not None. + forces: Forces [N]. Shape: (len(env_ids), len(body_ids), 3). Defaults to None. + torques: Torques [N·m]. Shape: (len(env_ids), len(body_ids), 3). Defaults to None. + positions: The positions [m] at which forces act. If `is_global` is True, these are global + positions expressed in the world frame. If `is_global` is False, these are offsets from the + body's CoM expressed in the body frame. If None, forces are assumed to act at the body's + CoM, independent of the `is_global` flag. + Shape: (len(env_ids), len(body_ids), 3). Defaults to None. + body_ids: Body indices. Defaults to None (all bodies). + env_ids: Environment indices. Defaults to None (all environments). + is_global: Whether the forces and torques are expressed in the global world frame or the local body frame. + Defaults to False. """ - # Resolve all indices - if (env_ids is None) or (env_ids == slice(None)): - env_ids = self._ALL_ENV_INDICES - if isinstance(env_ids, list): - env_ids = wp.array(env_ids, dtype=wp.int32, device=self.device) - if (body_ids is None) or (body_ids == slice(None)): - body_ids = self._ALL_BODY_INDICES - if isinstance(body_ids, list): - body_ids = wp.array(body_ids, dtype=wp.int32, device=self.device) + env_ids = self._resolve_env_ids(env_ids) + body_ids = self._resolve_body_ids(body_ids) if forces is None and torques is None: warnings.warn( - "No forces or torques provided. No force will be added.", + "No forces or torques provided. No force will be set.", UserWarning, stacklevel=2, ) return - # Get the link poses - if not self._link_poses_updated: - self._link_poses = self._get_link_pose_fn() - self._link_poses_updated = True - # Set the active flag to true + # Clear input buffers for the targeted environments before writing + self.reset(env_ids=env_ids) + self._active = True + self._dirty = True wp.launch( - set_forces_and_torques_at_position_index, + set_forces_to_dual_buffers_index, dim=(env_ids.shape[0], body_ids.shape[0]), inputs=[ env_ids, @@ -236,13 +325,13 @@ def set_forces_and_torques_index( forces, torques, positions, - self._link_poses, + self._global_force_w, + self._global_torque_w, + self._global_force_at_com_w, + self._local_force_b, + self._local_torque_b, is_global, ], - outputs=[ - self._composed_force_b, - self._composed_torque_b, - ], device=self.device, ) @@ -255,30 +344,23 @@ def add_forces_and_torques_mask( env_mask: wp.array | torch.Tensor | None = None, is_global: bool = False, ): - """Add forces and torques to the composed force and torque. - - Composed force and torque are the sum of all the forces and torques applied to the body. - It can compose global wrenches and local wrenches. The result is always in the link frame of the body. + """Add forces and torques into the input buffers using mask-based selection. - The user can provide any combination of forces, torques, and positions. - - .. note:: Users may want to call `reset` function after every simulation step to ensure no force is carried - over to the next step. However, this may not necessary if the user calls `set_forces_and_torques` function - instead of `add_forces_and_torques`. + Accumulates onto whatever is already in the buffers. Args: - forces: Forces. (num_envs, num_bodies, 3). Defaults to None. - torques: Torques. (num_envs, num_bodies, 3). Defaults to None. - positions: Positions. (num_envs, num_bodies, 3). Defaults to None. - body_mask: Body mask. (num_bodies). Defaults to None (all bodies). - env_mask: Environment mask. (num_envs). Defaults to None (all environments). - is_global: Whether the forces and torques are applied in the global frame. Defaults to False. - - Raises: - ValueError: If the type of the input is not supported. - ValueError: If the input is a slice and it is not None. + forces: Forces [N]. Shape: (num_envs, num_bodies, 3). Defaults to None. + torques: Torques [N·m]. Shape: (num_envs, num_bodies, 3). Defaults to None. + positions: The positions [m] at which forces act. If `is_global` is True, these are global + positions expressed in the world frame. If `is_global` is False, these are offsets from the + body's CoM expressed in the body frame. If None, forces are assumed to act at the body's + CoM, independent of the `is_global` flag. + Shape: (num_envs, num_bodies, 3). Defaults to None. + body_mask: Body mask. Shape: (num_bodies,). Defaults to None (all bodies). + env_mask: Environment mask. Shape: (num_envs,). Defaults to None (all environments). + is_global: Whether the forces and torques are expressed in the global world frame or the local body frame. + Defaults to False. """ - # Resolve all indices if env_mask is None: env_mask = self._ALL_ENV_MASK if body_mask is None: @@ -290,16 +372,12 @@ def add_forces_and_torques_mask( stacklevel=2, ) return - # Get the link poses - if not self._link_poses_updated: - self._link_poses = self._get_link_pose_fn() - self._link_poses_updated = True - # Set the active flag to true self._active = True + self._dirty = True wp.launch( - add_forces_and_torques_at_position_mask, + add_forces_to_dual_buffers_mask, dim=(self.num_envs, self.num_bodies), inputs=[ env_mask, @@ -307,13 +385,13 @@ def add_forces_and_torques_mask( forces, torques, positions, - self._link_poses, + self._global_force_w, + self._global_torque_w, + self._global_force_at_com_w, + self._local_force_b, + self._local_torque_b, is_global, ], - outputs=[ - self._composed_force_b, - self._composed_torque_b, - ], device=self.device, ) @@ -326,47 +404,45 @@ def set_forces_and_torques_mask( env_mask: wp.array | torch.Tensor | None = None, is_global: bool = False, ): - """Set forces and torques to the composed force and torque. + """Set forces and torques into the input buffers using mask-based selection. - Composed force and torque are the sum of all the forces and torques applied to the body. - It can compose global wrenches and local wrenches. The result is always in the link frame of the body. - - The user can provide any combination of forces, torques, and positions. + Resets the masked environments first, then writes the new values. This replaces any + previously accumulated forces/torques for the masked environments while leaving other + environments untouched. Args: - forces: Forces. (num_envs, num_bodies, 3). Defaults to None. - torques: Torques. (num_envs, num_bodies, 3). Defaults to None. - positions: Positions. (num_envs, num_bodies, 3). Defaults to None. - body_mask: Body mask. (num_bodies). Defaults to None (all bodies). - env_mask: Environment mask. (num_envs). Defaults to None (all environments). - is_global: Whether the forces and torques are applied in the global frame. Defaults to False. - - Raises: - ValueError: If the type of the input is not supported. - ValueError: If the input is a slice and it is not None. + forces: Forces [N]. Shape: (num_envs, num_bodies, 3). Defaults to None. + torques: Torques [N·m]. Shape: (num_envs, num_bodies, 3). Defaults to None. + positions: The positions [m] at which forces act. If `is_global` is True, these are global + positions expressed in the world frame. If `is_global` is False, these are offsets from the + body's CoM expressed in the body frame. If None, forces are assumed to act at the body's + CoM, independent of the `is_global` flag. + Shape: (num_envs, num_bodies, 3). Defaults to None. + body_mask: Body mask. Shape: (num_bodies,). Defaults to None (all bodies). + env_mask: Environment mask. Shape: (num_envs,). Defaults to None (all environments). + is_global: Whether the forces and torques are expressed in the global world frame or the local body frame. + Defaults to False. """ - # Resolve all indices if env_mask is None: env_mask = self._ALL_ENV_MASK if body_mask is None: body_mask = self._ALL_BODY_MASK if forces is None and torques is None: warnings.warn( - "No forces or torques provided. No force will be added.", + "No forces or torques provided. No force will be set.", UserWarning, stacklevel=2, ) return - # Get the link poses - if not self._link_poses_updated: - self._link_poses = self._get_link_pose_fn() - self._link_poses_updated = True - # Set the active flag to true + # Clear input buffers for the masked environments before writing + self.reset(env_mask=env_mask) + self._active = True + self._dirty = True wp.launch( - set_forces_and_torques_at_position_mask, + set_forces_to_dual_buffers_mask, dim=(self.num_envs, self.num_bodies), inputs=[ env_mask, @@ -374,74 +450,159 @@ def set_forces_and_torques_mask( forces, torques, positions, - self._link_poses, + self._global_force_w, + self._global_torque_w, + self._global_force_at_com_w, + self._local_force_b, + self._local_torque_b, is_global, ], - outputs=[ - self._composed_force_b, - self._composed_torque_b, + device=self.device, + ) + + def add_raw_buffers_from(self, other: WrenchComposer): + """Add another composer's raw input buffers into this composer's input buffers. + + This performs element-wise addition of all five input buffers from ``other`` into ``self``. + Useful for combining wrenches from multiple sources before composition. + + Args: + other: Another WrenchComposer whose input buffers will be added into this one. + """ + if not other._active: + return + if __debug__: + if other.num_envs != self.num_envs or other.num_bodies != self.num_bodies: + raise ValueError( + f"Cannot add buffers from composer with shape ({other.num_envs}, {other.num_bodies}) " + f"into composer with shape ({self.num_envs}, {self.num_bodies})." + ) + + self._active = True + self._dirty = True + + wp.launch( + add_raw_wrench_buffers, + dim=(self.num_envs, self.num_bodies), + inputs=[ + other._global_force_w, + other._global_torque_w, + other._global_force_at_com_w, + other._local_force_b, + other._local_torque_b, + self._global_force_w, + self._global_torque_w, + self._global_force_at_com_w, + self._local_force_b, + self._local_torque_b, ], device=self.device, ) - def reset(self, env_ids: wp.array | torch.Tensor | None = None, env_mask: wp.array | None = None): - """Reset the composed force and torque. + def compose_to_body_frame(self): + """Compose the five input buffers into the two output buffers in body frame. - This function will reset the composed force and torque to zero. - It will also make sure the link positions and quaternions are updated in the next call of the - `add_forces_and_torques` or `set_forces_and_torques` functions. + This corrects world-frame torques for the body's CoM position, rotates global forces and torques into the + body frame, then adds local-frame contributions. After this call, ``out_force_b`` and ``out_torque_b`` + contain the final composed wrench. - .. note:: This function should be called after every simulation step / reset to ensure no force is carried - over to the next step. + The dirty flag is cleared after composition. + """ + com_pos_w = self._get_com_pos_fn() + link_quat_w = self._get_link_quat_fn() - .. caution:: If both :attr:`env_ids` and :attr:`env_mask` are provided, then :attr:`env_mask` takes precedence - over :attr:`env_ids`. + wp.launch( + compose_wrench_to_body_frame, + dim=(self.num_envs, self.num_bodies), + inputs=[ + self._global_force_w, + self._global_torque_w, + self._global_force_at_com_w, + self._local_force_b, + self._local_torque_b, + com_pos_w, + link_quat_w, + self._out_force_b, + self._out_torque_b, + ], + device=self.device, + ) + self._dirty = False + + def reset( + self, + env_ids: wp.array | torch.Tensor | Sequence[int] | slice | None = None, + env_mask: wp.array | None = None, + ): + """Reset the wrench composer buffers. + + With no arguments, zeros all seven buffers (5 input + 2 output) and clears all flags. + With ``env_ids`` or ``env_mask``, performs a partial reset on the specified environments + using the reset kernels. + + .. caution:: If both ``env_ids`` and ``env_mask`` are provided, ``env_mask`` takes precedence. Args: env_ids: Environment indices. Defaults to None (all environments). env_mask: Environment mask. Defaults to None (all environments). """ if env_ids is None and env_mask is None: - self._composed_force_b.zero_() - self._composed_torque_b.zero_() + # Full reset: zero all 7 buffers + self._global_force_w.zero_() + self._global_torque_w.zero_() + self._global_force_at_com_w.zero_() + self._local_force_b.zero_() + self._local_torque_b.zero_() + self._out_force_b.zero_() + self._out_torque_b.zero_() self._active = False + self._dirty = False elif env_mask is not None: wp.launch( reset_wrench_composer_mask, dim=(self.num_envs, self.num_bodies), inputs=[ env_mask, - ], - outputs=[ - self._composed_force_b, - self._composed_torque_b, + self._global_force_w, + self._global_torque_w, + self._global_force_at_com_w, + self._local_force_b, + self._local_torque_b, + self._out_force_b, + self._out_torque_b, ], device=self.device, ) + self._dirty = True else: + # Partial reset via index if env_ids is None or env_ids == slice(None): env_ids = self._ALL_ENV_INDICES elif isinstance(env_ids, list): env_ids = wp.array(env_ids, dtype=wp.int32, device=self.device) elif isinstance(env_ids, torch.Tensor): env_ids = wp.from_torch(env_ids.to(torch.int32), dtype=wp.int32) + wp.launch( reset_wrench_composer_index, dim=(env_ids.shape[0], self.num_bodies), inputs=[ env_ids, - ], - outputs=[ - self._composed_force_b, - self._composed_torque_b, + self._global_force_w, + self._global_torque_w, + self._global_force_at_com_w, + self._local_force_b, + self._local_torque_b, + self._out_force_b, + self._out_torque_b, ], device=self.device, ) - self._link_poses_updated = False + self._dirty = True - """ - Deprecated functions. - """ + # ------------------------------------------------------------------ + # Deprecated methods + # ------------------------------------------------------------------ def add_forces_and_torques( self, @@ -452,10 +613,13 @@ def add_forces_and_torques( env_ids: torch.Tensor | None = None, is_global: bool = False, ): - """Deprecated, same as :meth:`add_forces_and_torques_index`.""" + """Deprecated, same as :meth:`add_forces_and_torques_index`. + + .. deprecated:: 4.5.33 + Use :meth:`add_forces_and_torques_index` instead. + """ warnings.warn( - "The function 'add_forces_and_torques' will be deprecated in a future release. Please" - " use 'add_forces_and_torques_index' instead.", + "The function 'add_forces_and_torques' is deprecated. Please use 'add_forces_and_torques_index' instead.", DeprecationWarning, stacklevel=2, ) @@ -470,11 +634,80 @@ def set_forces_and_torques( env_ids: wp.array | torch.Tensor | None = None, is_global: bool = False, ): - """Deprecated, same as :meth:`set_forces_and_torques_index`.""" + """Deprecated, same as :meth:`set_forces_and_torques_index`. + + .. deprecated:: 4.5.33 + Use :meth:`set_forces_and_torques_index` instead. + """ warnings.warn( - "The function 'set_forces_and_torques' will be deprecated in a future release. Please" - " use 'set_forces_and_torques_index' instead.", + "The function 'set_forces_and_torques' is deprecated. Please use 'set_forces_and_torques_index' instead.", DeprecationWarning, stacklevel=2, ) self.set_forces_and_torques_index(forces, torques, positions, body_ids, env_ids, is_global) + + # ------------------------------------------------------------------ + # Internal helpers + # ------------------------------------------------------------------ + + def _resolve_env_ids(self, env_ids: wp.array | torch.Tensor | list | slice | None) -> wp.array: + """Resolve environment IDs to a warp int32 array. + + Args: + env_ids: Environment indices as any supported type, or None for all environments. + + Returns: + Warp array of int32 environment indices. + + Raises: + TypeError: If ``env_ids`` is an unsupported type. + """ + if env_ids is None: + return self._ALL_ENV_INDICES + # Check tensor types before slice comparison (tensor == slice crashes) + if isinstance(env_ids, torch.Tensor): + if env_ids.dtype == torch.int64: + env_ids = env_ids.to(torch.int32) + return wp.from_torch(env_ids.contiguous(), dtype=wp.int32) + if isinstance(env_ids, wp.array): + return env_ids + if env_ids == slice(None): + return self._ALL_ENV_INDICES + if isinstance(env_ids, list): + return wp.array(env_ids, dtype=wp.int32, device=self.device) + raise TypeError( + f"env_ids must be None, slice(None), list, torch.Tensor, or wp.array, got {type(env_ids).__name__}" + ) + + def _resolve_body_ids(self, body_ids: wp.array | torch.Tensor | list | slice | None) -> wp.array: + """Resolve body IDs to a warp int32 array. + + Args: + body_ids: Body indices as any supported type, or None for all bodies. + + Returns: + Warp array of int32 body indices. + + Raises: + TypeError: If ``body_ids`` is an unsupported type. + """ + if body_ids is None: + return self._ALL_BODY_INDICES + if isinstance(body_ids, torch.Tensor): + if body_ids.dtype == torch.int64: + body_ids = body_ids.to(torch.int32) + return wp.from_torch(body_ids.contiguous(), dtype=wp.int32) + if isinstance(body_ids, wp.array): + return body_ids + if body_ids == slice(None): + return self._ALL_BODY_INDICES + if isinstance(body_ids, list): + return wp.array(body_ids, dtype=wp.int32, device=self.device) + raise TypeError( + f"body_ids must be None, slice(None), list, torch.Tensor, or wp.array, got {type(body_ids).__name__}" + ) + + def _ensure_composed(self): + """Compose input buffers into output buffers if dirty.""" + if self._dirty: + self.compose_to_body_frame() diff --git a/source/isaaclab/isaaclab/visualizers/base_visualizer.py b/source/isaaclab/isaaclab/visualizers/base_visualizer.py index 2480bc89bbaa..b0fc5a81088f 100644 --- a/source/isaaclab/isaaclab/visualizers/base_visualizer.py +++ b/source/isaaclab/isaaclab/visualizers/base_visualizer.py @@ -136,7 +136,7 @@ def get_visualized_env_ids(self) -> list[int] | None: Returns: Visualized environment ids, or ``None`` for all environments. """ - return getattr(self, "_env_ids", None) + return self._env_ids def _compute_visualized_env_ids(self) -> list[int] | None: """Compute which environment indices to visualize from config. @@ -146,28 +146,21 @@ def _compute_visualized_env_ids(self) -> list[int] | None: """ if self._scene_data_provider is None: return None - filter_mode = getattr(self.cfg, "env_filter_mode", "none") - if filter_mode == "none": - return None - + cfg = self.cfg num_envs = self._scene_data_provider.get_metadata().get("num_envs", 0) if num_envs <= 0: - logger.debug("[Visualizer] num_envs is 0 or missing from provider metadata; env filtering disabled.") + logger.debug("[Visualizer] num_envs is 0 or missing from provider metadata; env selection disabled.") return None - if filter_mode == "env_ids": - env_ids_cfg = getattr(self.cfg, "env_filter_ids", None) - if env_ids_cfg is not None and len(env_ids_cfg) > 0: - return [i for i in env_ids_cfg if 0 <= i < num_envs] - return None - if filter_mode == "random_n": - count = int(getattr(self.cfg, "env_filter_random_n", 0)) - if count <= 0: - return None - count = min(count, num_envs) - seed = int(getattr(self.cfg, "env_filter_seed", 0)) - rng = random.Random(seed) - return sorted(rng.sample(range(num_envs), count)) - logger.warning("[Visualizer] Unknown env_filter_mode='%s'; defaulting to all envs.", filter_mode) + # Explicit list wins; never combine with random cap-only mode. + if cfg.visible_env_indices is not None: + return [i for i in cfg.visible_env_indices if 0 <= i < num_envs] + + max_visible = getattr(cfg, "max_visible_envs", None) + # Random subset only for cap-only mode: needs a cap and no explicit indices (see VisualizerCfg). + if max_visible is not None and getattr(cfg, "randomly_sample_visible_envs", True) and int(max_visible) >= 0: + k = min(int(max_visible), num_envs) + # k == 0: sample(range(n), 0) is []; contiguous resolver used the same convention. + return sorted(random.sample(range(num_envs), k)) return None def get_rendering_dt(self) -> float | None: @@ -187,6 +180,14 @@ def set_camera_view(self, eye: tuple, target: tuple) -> None: """ pass + def _resolve_cfg_camera_pose( + self, _visualizer_name: str + ) -> tuple[tuple[float, float, float], tuple[float, float, float]]: + """Resolve camera pose from cfg eye/lookat fields.""" + eye = tuple(float(v) for v in self.cfg.eye) + lookat = tuple(float(v) for v in self.cfg.lookat) + return eye, lookat + def _resolve_camera_pose_from_usd_path( self, usd_path: str ) -> tuple[tuple[float, float, float], tuple[float, float, float]] | None: diff --git a/source/isaaclab/isaaclab/visualizers/visualizer_cfg.py b/source/isaaclab/isaaclab/visualizers/visualizer_cfg.py index a96e3c04d2b5..1ee4cde038b5 100644 --- a/source/isaaclab/isaaclab/visualizers/visualizer_cfg.py +++ b/source/isaaclab/isaaclab/visualizers/visualizer_cfg.py @@ -34,32 +34,33 @@ class VisualizerCfg: enable_live_plots: bool = True """Enable live plotting of data.""" - camera_position: tuple[float, float, float] = (8.0, 8.0, 3.0) - """Initial camera position (x, y, z) in world coordinates.""" + eye: tuple[float, float, float] = (7.5, 7.5, 7.5) + """Initial camera eye position (x, y, z) in world coordinates.""" - camera_target: tuple[float, float, float] = (0.0, 0.0, 0.0) - """Initial camera target/look-at point (x, y, z) in world coordinates.""" + lookat: tuple[float, float, float] = (0.0, 0.0, 0.0) + """Initial camera look-at point (x, y, z) in world coordinates.""" - camera_source: Literal["cfg", "usd_path"] = "cfg" - """Camera source mode: 'cfg' uses camera_position/target, 'usd_path' follows a USD camera prim.""" + cam_source: Literal["cfg", "prim_path"] = "cfg" + """Camera source mode: 'cfg' uses eye/lookat, 'prim_path' follows a camera prim.""" - camera_usd_path: str = "/World/envs/env_0/Camera" - """Absolute USD path to a camera prim when camera_source='usd_path'.""" + cam_prim_path: str = "/World/envs/env_0/Camera" + """Absolute USD path to a camera prim when cam_source='prim_path'.""" - env_filter_mode: Literal["none", "env_ids", "random_n"] = "none" - """Env filter mode: 'none', 'env_ids', or 'random_n'.""" + max_visible_envs: int | None = None + """Upper bound on how many envs are shown. - env_filter_random_n: int = 64 - """If env_filter_mode='random_n', number of envs to sample.""" + * If visible_env_indices is not None, then this field will apply also + to the explicit env indices set to the visible_env_indices. + """ - env_filter_seed: int = 0 - """Seed for deterministic env sampling.""" + visible_env_indices: list[int] | None = None + """env indices to visualize in order (out-of-range indices are dropped).""" - env_filter_ids: list[int] = [i for i in range(0, 64, 4)] - """If env_filter_mode='env_ids', only these env indices are shown. + randomly_sample_visible_envs: bool = True + """If ``max_visible_envs`` is provided, when enabled, selected visible envs are randomly sampled. + If disabled, the first ``max_visible_envs`` envs are selected. - This improves performance, particularly for large-scale training, by reducing scene updates sent to visualizers. - Note, OV visualizer only applies a cosmetic visibility toggle (no performance gain). + * Note: ``visible_env_indices`` overrides this field. """ def get_visualizer_type(self) -> str | None: diff --git a/source/isaaclab/setup.py b/source/isaaclab/setup.py index db3d42f1f279..7d14504f48a3 100644 --- a/source/isaaclab/setup.py +++ b/source/isaaclab/setup.py @@ -30,8 +30,8 @@ # procedural-generation "trimesh", "pyglet>=2.1.6,<3", - "mujoco==3.5.0", - "mujoco-warp==3.5.0.2", + "mujoco==3.6.0", + "mujoco-warp==3.6.0", # image processing "transformers==4.57.6", "einops", # needed for transformers, doesn't always auto-install diff --git a/source/isaaclab/test/app/test_kwarg_launch.py b/source/isaaclab/test/app/test_kwarg_launch.py index 0dffa89764a0..25fe56b69232 100644 --- a/source/isaaclab/test/app/test_kwarg_launch.py +++ b/source/isaaclab/test/app/test_kwarg_launch.py @@ -43,25 +43,27 @@ def test_set_visualizer_settings_stores_values(monkeypatch: pytest.MonkeyPatch): monkeypatch.setattr(app_launcher_module, "get_settings_manager", lambda: settings) launcher = AppLauncher.__new__(AppLauncher) - launcher._set_visualizer_settings({"visualizer": ["viser", "rerun"], "visualizer_max_worlds": 0}) + launcher._set_visualizer_settings({"visualizer": ["viser", "rerun"], "max_visible_envs": 0}) assert settings.values == { "/isaaclab/visualizer/types": "viser rerun", "/isaaclab/visualizer/explicit": False, "/isaaclab/visualizer/disable_all": False, - "/isaaclab/visualizer/max_worlds": 0, + "/isaaclab/visualizer/max_visible_envs": 0, } -def test_set_visualizer_settings_rejects_negative_max_worlds(monkeypatch: pytest.MonkeyPatch): +def test_set_visualizer_settings_rejects_negative_max_visible_envs( + monkeypatch: pytest.MonkeyPatch, +): def _unexpected_settings_manager(): raise AssertionError("settings manager should not be queried for invalid values") monkeypatch.setattr(app_launcher_module, "get_settings_manager", _unexpected_settings_manager) launcher = AppLauncher.__new__(AppLauncher) - with pytest.raises(ValueError, match="Invalid value for --visualizer_max_worlds: -5"): - launcher._set_visualizer_settings({"visualizer": ["viser"], "visualizer_max_worlds": -5}) + with pytest.raises(ValueError, match="Invalid value for --max_visible_envs: -5"): + launcher._set_visualizer_settings({"visualizer": ["viser"], "max_visible_envs": -5}) def test_set_visualizer_settings_suppresses_settings_manager_errors(monkeypatch: pytest.MonkeyPatch): @@ -71,17 +73,17 @@ def _raise_settings_error(): monkeypatch.setattr(app_launcher_module, "get_settings_manager", _raise_settings_error) launcher = AppLauncher.__new__(AppLauncher) - launcher._set_visualizer_settings({"visualizer": ["viser"], "visualizer_max_worlds": 3}) + launcher._set_visualizer_settings({"visualizer": ["viser"], "max_visible_envs": 3}) def test_parse_visualizer_csv_accepts_comma_delimited_values(): - parsed = app_launcher_module._parse_visualizer_csv("kit,newton,rerun,viser") + parsed = app_launcher_module.AppLauncher._parse_visualizer_csv("kit,newton,rerun,viser") assert parsed == ["kit", "newton", "rerun", "viser"] def test_parse_visualizer_csv_rejects_spaces_between_entries(): with pytest.raises(argparse.ArgumentTypeError, match="spaces are not allowed"): - app_launcher_module._parse_visualizer_csv("kit, newton") + app_launcher_module.AppLauncher._parse_visualizer_csv("kit, newton") def test_resolve_visualizer_settings_rejects_none_with_others(): diff --git a/source/isaaclab/test/assets/test_articulation_iface.py b/source/isaaclab/test/assets/test_articulation_iface.py index 14af782c2a5b..cbec2781065b 100644 --- a/source/isaaclab/test/assets/test_articulation_iface.py +++ b/source/isaaclab/test/assets/test_articulation_iface.py @@ -602,6 +602,72 @@ def test_find_joints_single(self, backend, num_instances, num_joints, num_bodies assert names == [first_joint] +# --------------------------------------------------------------------------- +# Tests: resolve_matching_names caching behavior +# --------------------------------------------------------------------------- + + +_non_mock_backends = pytest.mark.parametrize("backend", [b for b in BACKENDS if b != "mock"], indirect=False) + + +class TestResolveMatchingNamesCache: + """Test that resolve_matching_names caching returns correct, isolated results.""" + + @_non_mock_backends + @pytest.mark.parametrize("num_instances, num_joints, num_bodies", [(2, 6, 7)]) + @_default_devices + def test_unmatched_regex_raises(self, backend, num_instances, num_joints, num_bodies, device): + """ValueError from resolve_matching_names propagates correctly.""" + art, _ = get_articulation(backend, num_instances, num_joints, num_bodies, device=device) + with pytest.raises(ValueError): + art.find_bodies("nonexistent_body_xyz") + with pytest.raises(ValueError): + art.find_joints("nonexistent_joint_xyz") + + @_backends + @pytest.mark.parametrize("num_instances, num_joints, num_bodies", [(2, 6, 7)]) + @_default_devices + def test_mutating_result_does_not_corrupt_cache( + self, backend, num_instances, num_joints, num_bodies, device, articulation_iface + ): + """Mutating returned lists must not affect future cached results.""" + art, _ = articulation_iface + + for finder, expected_len in [("find_bodies", num_bodies), ("find_joints", num_joints)]: + idx1, names1 = getattr(art, finder)(".*") + assert len(idx1) == expected_len + + idx1.clear() + names1.append("corrupted") + + idx2, names2 = getattr(art, finder)(".*") + assert len(idx2) == expected_len + assert "corrupted" not in names2 + + @_non_mock_backends + @pytest.mark.parametrize("num_instances, num_joints, num_bodies", [(2, 6, 7)]) + @_default_devices + def test_find_with_multiple_patterns(self, backend, num_instances, num_joints, num_bodies, device): + """Passing a list of regex patterns works correctly.""" + art, _ = get_articulation(backend, num_instances, num_joints, num_bodies, device=device) + idx, names = art.find_joints(["joint_0", "joint_1"]) + assert "joint_0" in names + assert "joint_1" in names + assert len(names) == 2 + + @_non_mock_backends + @pytest.mark.parametrize("num_instances, num_joints, num_bodies", [(2, 6, 7)]) + @_default_devices + def test_find_with_preserve_order(self, backend, num_instances, num_joints, num_bodies, device): + """preserve_order=True returns names in the order of the input patterns.""" + art, _ = get_articulation(backend, num_instances, num_joints, num_bodies, device=device) + idx_fwd, names_fwd = art.find_joints(["joint_1", "joint_0"], preserve_order=True) + assert names_fwd == ["joint_1", "joint_0"] + + idx_rev, names_rev = art.find_joints(["joint_0", "joint_1"], preserve_order=True) + assert names_rev == ["joint_0", "joint_1"] + + # --------------------------------------------------------------------------- # Tests: ArticulationData root state properties # --------------------------------------------------------------------------- diff --git a/source/isaaclab/test/install_ci/conftest.py b/source/isaaclab/test/install_ci/conftest.py index 226af44a65f6..c4bbc94ab200 100644 --- a/source/isaaclab/test/install_ci/conftest.py +++ b/source/isaaclab/test/install_ci/conftest.py @@ -20,6 +20,8 @@ _CYAN_BRIGHT = "\033[96m" _RESET = "\033[0m" +_EXECUTION_ENVIRONMENT_KEY = pytest.StashKey[_utils.ExecutionEnvironment]() + # Fixtures @@ -79,16 +81,26 @@ def wheel_path() -> Path | None: def pytest_configure(config: pytest.Config) -> None: config.addinivalue_line("markers", "bug: bug-regression tests (use bug id as argument)") config.addinivalue_line("markers", "gpu: tests that require a GPU") - config.addinivalue_line("markers", "docker_only: tests that only run inside Docker") - config.addinivalue_line("markers", "needs_network: tests that require network access") + config.addinivalue_line("markers", "docker: tests that only run inside Docker") + config.addinivalue_line("markers", "native: tests that only run natively (not in Docker)") config.addinivalue_line("markers", "slow: tests that take a long time") config.addinivalue_line("markers", "uv: tests that require the uv package manager") + try: + config.stash[_EXECUTION_ENVIRONMENT_KEY] = _utils.detect_execution_environment() + except ValueError as exc: + raise pytest.UsageError(str(exc)) from exc + # Enable real-time output when pytest capture is disabled (-s) capture = config.getoption("capture", default="fd") _utils.stream_output = capture == "no" +def pytest_report_header(config: pytest.Config) -> str: + """Show the detected install_ci execution environment in the test header.""" + return f"install_ci execution environment: {config.stash[_EXECUTION_ENVIRONMENT_KEY]}" + + def pytest_runtest_logreport(report: pytest.TestReport) -> None: """Print a newline after the PASSED/FAILED/SKIPPED result.""" if report.when == "call" or (report.when == "setup" and report.skipped): @@ -97,11 +109,12 @@ def pytest_runtest_logreport(report: pytest.TestReport) -> None: @pytest.hookimpl(tryfirst=True) def pytest_collection_modifyitems(config: pytest.Config, items: list[pytest.Item]) -> None: - """Dynamically map marker arguments from `@pytest.mark.bug("...")` to standalone markers. + """Map dynamic bug markers and skip items with mismatched env markers. This allows filtering by bug ID natively in pytest: `-m "nvbugs_5968136"` instead of the (unsupported natively) `-m "bug('nvbugs_5968136')"`. """ + execution_environment = config.stash[_EXECUTION_ENVIRONMENT_KEY] known_bugs = set() for item in items: for mark in item.iter_markers(name="bug"): @@ -117,3 +130,12 @@ def pytest_collection_modifyitems(config: pytest.Config, items: list[pytest.Item for arg in mark.args: if isinstance(arg, str): item.add_marker(arg) + + marker_names = {mark.name for mark in item.iter_markers()} + try: + skip_reason = _utils.get_execution_environment_skip_reason(marker_names, execution_environment) + except ValueError as exc: + raise pytest.UsageError(f"{item.nodeid}: {exc}") from exc + + if skip_reason: + item.add_marker(pytest.mark.skip(reason=skip_reason)) diff --git a/source/isaaclab/test/install_ci/pytest.ini b/source/isaaclab/test/install_ci/pytest.ini index 26abb3f86add..67eac9b17ca7 100644 --- a/source/isaaclab/test/install_ci/pytest.ini +++ b/source/isaaclab/test/install_ci/pytest.ini @@ -8,8 +8,8 @@ python_files = markers = bug: bug-regression tests (use bug id as argument) gpu: tests that require a GPU - docker_only: tests that only run inside Docker - needs_network: tests that require network access + docker: tests that only run inside Docker + native: tests that only run natively (not in Docker) slow: tests that take a long time uv: tests that require the uv package manager timeout = 1200 diff --git a/source/isaaclab/test/install_ci/test_environment_markers.py b/source/isaaclab/test/install_ci/test_environment_markers.py new file mode 100644 index 000000000000..c002353b5c3e --- /dev/null +++ b/source/isaaclab/test/install_ci/test_environment_markers.py @@ -0,0 +1,77 @@ +# 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 + +"""Unit tests for install_ci execution-environment marker handling.""" + +from __future__ import annotations + +import pytest +from utils import ( + detect_execution_environment, + get_execution_environment_skip_reason, +) + + +class TestDetectExecutionEnvironment: + """Tests for detect_execution_environment().""" + + def test_uses_override(self, tmp_path): + environment = detect_execution_environment( + environ={"ISAACLAB_INSTALL_CI_ENV": "docker"}, + filesystem_root=tmp_path, + ) + + assert environment == "docker" + + def test_detects_marker_file(self, tmp_path): + (tmp_path / ".dockerenv").touch() + + environment = detect_execution_environment(environ={}, filesystem_root=tmp_path) + + assert environment == "docker" + + def test_detects_cgroup_hint(self, tmp_path): + cgroup_path = tmp_path / "proc" / "self" + cgroup_path.mkdir(parents=True) + (cgroup_path / "cgroup").write_text("0::/docker/container-id\n", encoding="utf-8") + + environment = detect_execution_environment(environ={}, filesystem_root=tmp_path) + + assert environment == "docker" + + def test_defaults_to_native(self, tmp_path): + environment = detect_execution_environment(environ={}, filesystem_root=tmp_path) + + assert environment == "native" + + def test_rejects_invalid_override(self, tmp_path): + with pytest.raises(ValueError, match="ISAACLAB_INSTALL_CI_ENV"): + detect_execution_environment( + environ={"ISAACLAB_INSTALL_CI_ENV": "virtual-machine"}, + filesystem_root=tmp_path, + ) + + +class TestGetExecutionEnvironmentSkipReason: + """Tests for get_execution_environment_skip_reason().""" + + @pytest.mark.parametrize( + ("marker_names", "execution_environment", "expected_reason"), + [ + ({"docker"}, "native", "requires Docker execution environment, detected native"), + ({"native"}, "docker", "requires native execution environment, detected docker"), + ({"docker"}, "docker", None), + ({"native"}, "native", None), + (set(), "native", None), + ], + ) + def test_skip_reason(self, marker_names, execution_environment, expected_reason): + skip_reason = get_execution_environment_skip_reason(marker_names, execution_environment) + + assert skip_reason == expected_reason + + def test_rejects_conflicting_markers(self): + with pytest.raises(ValueError, match="docker"): + get_execution_environment_skip_reason({"docker", "native"}, "native") diff --git a/source/isaaclab/test/install_ci/test_isaaclabx_i_newton.py b/source/isaaclab/test/install_ci/test_isaaclabx_i_newton.py new file mode 100644 index 000000000000..4f7740dfe6fd --- /dev/null +++ b/source/isaaclab/test/install_ci/test_isaaclabx_i_newton.py @@ -0,0 +1,61 @@ +# 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 + +"""Test installing isaaclab_newton and running its test suite.""" + +from __future__ import annotations + +import shutil + +import pytest +from utils import UV_Mixin, find_isaaclab_root + + +class Test_Install_Newton(UV_Mixin): + """Install ./isaaclab.sh -i newton and run the isaaclab_newton test suite.""" + + @classmethod + def setup_class(cls): + # check if uv is available + if not shutil.which("uv"): + pytest.skip("uv is not available") + + # check if isaacsim is importable + # or "_isaac_sim" link is present + try: + import isaacsim # noqa: F401 + except ImportError: + print("[DEBUG] Module isaacsim is not importable") + isaac_sim_link = find_isaaclab_root() / "_isaac_sim" + if not isaac_sim_link.exists(): + print(f'[DEBUG] Link "{isaac_sim_link}" does not exist') + pytest.skip("isaacsim is not importable and _isaac_sim link not found, skipping") + + @pytest.mark.uv + @pytest.mark.gpu + @pytest.mark.slow + @pytest.mark.native + @pytest.mark.timeout(3600) + def test_install_newton_and_run_tests(self, isaaclab_root): + """Install newton extension and run the isaaclab_newton test suite.""" + + try: + self.create_uv_env(isaaclab_root) + + # ./isaaclab.sh -i newton + result = self.run_in_uv_env([str(self.cli_script), "-i", "newton"], cwd=isaaclab_root) + assert result.returncode == 0, f"isaaclab -i newton failed:\n{result.stdout}\n{result.stderr}" + + # Run isaaclab_newton test suite + test_dir = str(isaaclab_root / "source" / "isaaclab_newton" / "test") + result = self.run_in_uv_env( + ["python", "-m", "pytest", test_dir, "-sv", "--tb=short"], + cwd=isaaclab_root, + ) + output = result.stdout + result.stderr + assert result.returncode == 0, f"isaaclab_newton tests failed (rc={result.returncode}):\n{output}" + + finally: + self.destroy_uv_env() diff --git a/source/isaaclab/test/install_ci/test_isaaclabx_i_physx.py b/source/isaaclab/test/install_ci/test_isaaclabx_i_physx.py new file mode 100644 index 000000000000..04bf2b346b23 --- /dev/null +++ b/source/isaaclab/test/install_ci/test_isaaclabx_i_physx.py @@ -0,0 +1,61 @@ +# 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 + +"""Test installing isaaclab_physx and running its test suite.""" + +from __future__ import annotations + +import shutil + +import pytest +from utils import UV_Mixin, find_isaaclab_root + + +class Test_Install_Physx(UV_Mixin): + """Install ./isaaclab.sh -i physx and run the isaaclab_physx test suite.""" + + @classmethod + def setup_class(cls): + # check if uv is available + if not shutil.which("uv"): + pytest.skip("uv is not available") + + # check if isaacsim is importable + # or "_isaac_sim" link is present + try: + import isaacsim # noqa: F401 + except ImportError: + print("[DEBUG] Module isaacsim is not importable") + isaac_sim_link = find_isaaclab_root() / "_isaac_sim" + if not isaac_sim_link.exists(): + print(f'[DEBUG] Link "{isaac_sim_link}" does not exist') + pytest.skip("isaacsim is not importable and _isaac_sim link not found, skipping") + + @pytest.mark.uv + @pytest.mark.gpu + @pytest.mark.slow + @pytest.mark.native + @pytest.mark.timeout(3600) + def test_install_physx_and_run_tests(self, isaaclab_root): + """Install physx extension and run the isaaclab_physx test suite.""" + + try: + self.create_uv_env(isaaclab_root) + + # ./isaaclab.sh -i physx + result = self.run_in_uv_env([str(self.cli_script), "-i", "physx"], cwd=isaaclab_root) + assert result.returncode == 0, f"isaaclab -i physx failed:\n{result.stdout}\n{result.stderr}" + + # Run isaaclab_physx test suite + test_dir = str(isaaclab_root / "source" / "isaaclab_physx" / "test") + result = self.run_in_uv_env( + ["python", "-m", "pytest", test_dir, "-sv", "--tb=short"], + cwd=isaaclab_root, + ) + output = result.stdout + result.stderr + assert result.returncode == 0, f"isaaclab_physx tests failed (rc={result.returncode}):\n{output}" + + finally: + self.destroy_uv_env() diff --git a/source/isaaclab/test/install_ci/test_isaaclabx_uv_smoke.py b/source/isaaclab/test/install_ci/test_isaaclabx_uv_smoke.py index d0fb0fee6b5d..8bff8426e1fe 100644 --- a/source/isaaclab/test/install_ci/test_isaaclabx_uv_smoke.py +++ b/source/isaaclab/test/install_ci/test_isaaclabx_uv_smoke.py @@ -55,8 +55,8 @@ def test_isaaclab_install_assets(self, isaaclab_root): @pytest.mark.uv @pytest.mark.timeout(300) - def test_isaaclab_newton_installs_isaaclab_physx(self, isaaclab_root): - """Run ./isaaclab.x -i 'newton' and verify isaaclab_physx is importable.""" + def test_isaaclab_newton_installs_isaaclab_newton(self, isaaclab_root): + """Run ./isaaclab.x -i 'newton' and verify isaaclab_newton is importable.""" try: self.create_uv_env(isaaclab_root) @@ -65,9 +65,9 @@ def test_isaaclab_newton_installs_isaaclab_physx(self, isaaclab_root): result = self.run_in_uv_env([str(self.cli_script), "-i", "newton"], cwd=isaaclab_root) assert result.returncode == 0, f"isaaclab -i newton failed:\n{result.stdout}\n{result.stderr}" - # import isaaclab_physx - result = self.run_in_uv_env(["python", "-c", "import isaaclab_physx; print(isaaclab_physx.__version__)"]) - assert result.returncode == 0, f"import isaaclab_physx failed:\n{result.stdout}\n{result.stderr}" + # import isaaclab_newton + result = self.run_in_uv_env(["python", "-c", "import isaaclab_newton; print(isaaclab_newton.__version__)"]) + assert result.returncode == 0, f"import isaaclab_newton failed:\n{result.stdout}\n{result.stderr}" finally: self.destroy_uv_env() diff --git a/source/isaaclab/test/install_ci/utils.py b/source/isaaclab/test/install_ci/utils.py index 08f046e5bb8c..85055f26c0bc 100644 --- a/source/isaaclab/test/install_ci/utils.py +++ b/source/isaaclab/test/install_ci/utils.py @@ -15,6 +15,7 @@ import sys import time from pathlib import Path +from typing import Literal _DIM = "\033[2m" _MAGENTA = "\033[95m" @@ -24,6 +25,73 @@ # Set to True by conftest.py when pytest runs with -s / --capture=no. stream_output: bool = False +# ISAACLAB_INSTALL_CI_ENV can be set to override execution +# environment detection in install_ci tests +# (for testing the testing while testing). + +ExecutionEnvironment = Literal["docker", "native"] + + +def detect_execution_environment( + environ: dict[str, str] | None = None, + filesystem_root: Path | None = None, +) -> ExecutionEnvironment: + """Detect whether install_ci tests are running in Docker or natively.""" + env = environ if environ is not None else os.environ + root = filesystem_root if filesystem_root is not None else Path("/") + + override = env.get("ISAACLAB_INSTALL_CI_ENV") + if override: + cleaned = override.strip().lower() + if cleaned not in ("docker", "native"): + raise ValueError(f"ISAACLAB_INSTALL_CI_ENV must be 'docker' or 'native', got: {override!r}") + return cleaned # type: ignore[return-value] + + if (root / ".dockerenv").exists() or (root / "run" / ".containerenv").exists(): + return "docker" + + for cgroup_path in (root / "proc" / "1" / "cgroup", root / "proc" / "self" / "cgroup"): + try: + cgroup_text = cgroup_path.read_text(encoding="utf-8", errors="ignore") + except OSError: + continue + if any( + hint in cgroup_text + for hint in ( + "docker", + "containerd", + "kubepods", + "libpod", + "podman", + ) + ): + return "docker" + + if env.get("container"): + return "docker" + + return "native" + + +def get_execution_environment_skip_reason( + marker_names: set[str], + execution_environment: ExecutionEnvironment, +) -> str | None: + """Return a skip reason when environment markers do not match the runtime.""" + has_docker = "docker" in marker_names + has_native = "native" in marker_names + + if has_docker and has_native: + raise ValueError("tests cannot be marked with both 'docker' and 'native'") + + if has_docker and execution_environment != "docker": + return f"requires Docker execution environment, detected {execution_environment}" + + if has_native and execution_environment != "native": + return f"requires native execution environment, detected {execution_environment}" + + return None + def find_isaaclab_root() -> Path: """Walk up from this file to find the repo root (contains isaaclab.sh).""" @@ -76,6 +144,7 @@ def run_cmd( stderr=subprocess.STDOUT, text=True, ) + assert proc.stdout is not None lines: list[str] = [] try: for line in proc.stdout: @@ -146,6 +215,9 @@ def create_uv_env(self, isaaclab_root: Path, env_name: str = "") -> None: assert result.returncode == 0, f"uv env creation failed:\n{result.stdout}\n{result.stderr}" assert self.env_path.exists(), f"Expected env directory {self.env_path} was not created" + # Prevent the venv from being tracked by git. + (self.env_path / ".gitignore").write_text("*\n") + self.python = (self.env_path / "Scripts" / "python.exe") if _IS_WINDOWS else (self.env_path / "bin" / "python") assert self.python.exists(), f"Python executable not found at {self.python}" 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_multi_mesh_ray_caster_camera.py b/source/isaaclab/test/sensors/test_multi_mesh_ray_caster_camera.py index 2b079760e16a..8657c938c691 100644 --- a/source/isaaclab/test/sensors/test_multi_mesh_ray_caster_camera.py +++ b/source/isaaclab/test/sensors/test_multi_mesh_ray_caster_camera.py @@ -777,3 +777,35 @@ def test_output_equal_to_usd_camera_when_intrinsics_set(setup_simulation): ) del camera_usd, camera_warp + + +@pytest.mark.isaacsim_ci +def test_image_mesh_ids_identifies_hit_mesh(setup_simulation): + """image_mesh_ids must contain 0 for ground-plane hits (only one mesh registered).""" + sim, dt, camera_cfg = setup_simulation + + cfg = copy.deepcopy(camera_cfg) + cfg.update_mesh_ids = True + cfg.data_types = ["distance_to_camera"] + + camera = MultiMeshRayCasterCamera(cfg=cfg) + sim.reset() + camera.update(dt) + + mesh_ids = camera.data.image_mesh_ids # shape (N, H, W, 1), dtype torch.int16 + assert mesh_ids is not None, "image_mesh_ids should not be None when update_mesh_ids=True" + assert mesh_ids.shape[-1] == 1 + assert mesh_ids.dtype == torch.int16 + + # Identify actual hits via distance < inf. This relies on depth_clipping_behavior="none" + # (the default), which leaves missed rays at the Warp-kernel fill value of inf. + # Under "max" clipping, missed rays would be clamped to a finite max_distance, making + # the inf comparison incorrect. + hit_mask = camera.data.output["distance_to_camera"][0, :, :, 0] < float("inf") + assert hit_mask.any(), "Expected at least some rays to hit the ground plane" + + # All hits against the single registered mesh must carry mesh_id=0 (first mesh index). + hit_mesh_ids = mesh_ids[0, :, :, 0][hit_mask] + assert torch.all(hit_mesh_ids == 0), ( + f"All hits against the single ground mesh must have mesh_id=0, got: {hit_mesh_ids.unique()}" + ) diff --git a/source/isaaclab/test/sensors/test_ray_caster.py b/source/isaaclab/test/sensors/test_ray_caster.py index 944287c549b4..4e29b25ce351 100644 --- a/source/isaaclab/test/sensors/test_ray_caster.py +++ b/source/isaaclab/test/sensors/test_ray_caster.py @@ -18,7 +18,9 @@ # Import after app launch import warp as wp -from isaaclab.utils.math import matrix_from_quat, quat_from_euler_xyz, random_orientation +from isaaclab.sensors.ray_caster.kernels import quat_yaw_only as _quat_yaw_only_func +from isaaclab.utils.math import matrix_from_quat, quat_from_euler_xyz, random_orientation, yaw_quat +from isaaclab.utils.warp.kernels import raycast_mesh_masked_kernel as _raycast_mesh_masked_kernel from isaaclab.utils.warp.ops import convert_to_warp_mesh, raycast_dynamic_meshes, raycast_mesh @@ -239,3 +241,283 @@ def test_raycast_random_cube(raycast_setup): torch.testing.assert_close(ray_distance, ray_distance_m) torch.testing.assert_close(ray_normal, ray_normal_m) 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) +# --------------------------------------------------------------------------- + +_SENTINEL = -2.0 # value pre-filled into output buffers; chosen outside [-1, 1] so it cannot +# equal any component of a unit-length surface normal, making "not written" assertions unambiguous. + + +def _make_masked_buffers(device, n_envs, n_rays): + """Allocate all warp buffers needed by raycast_mesh_masked_kernel. + + ray_dist_w and ray_normal_w are pre-filled with _SENTINEL so that tests can + meaningfully assert those buffers were *not* written when the corresponding + return flag is 0. + """ + ray_starts_w = wp.zeros((n_envs, n_rays), dtype=wp.vec3f, device=device) + ray_dirs_w = wp.zeros((n_envs, n_rays), dtype=wp.vec3f, device=device) + ray_hits_w = wp.zeros((n_envs, n_rays), dtype=wp.vec3f, device=device) + ray_dist_w = wp.zeros((n_envs, n_rays), dtype=wp.float32, device=device) + wp.to_torch(ray_dist_w).fill_(_SENTINEL) + ray_normal_w = wp.zeros((n_envs, n_rays), dtype=wp.vec3f, device=device) + wp.to_torch(ray_normal_w).fill_(_SENTINEL) + return ray_starts_w, ray_dirs_w, ray_hits_w, ray_dist_w, ray_normal_w + + +def test_raycast_mesh_masked_kernel_hits_only(raycast_setup): + """return_distance=0, return_normal=0: only ray_hits are written on a hit.""" + device = raycast_setup["device"] + mesh_id = raycast_setup["single_mesh_id"] + expected_hits = raycast_setup["expected_ray_hits"] # shape (1, 2, 3) + + n_envs, n_rays = 1, 2 + ray_starts_w, ray_dirs_w, ray_hits_w, ray_dist_w, ray_normal_w = _make_masked_buffers(device, n_envs, n_rays) + env_mask = wp.array([True], dtype=wp.bool, device=device) + + wp.to_torch(ray_starts_w)[:] = torch.tensor([[[0, -0.35, -5], [0.25, 0.35, -5]]], device=device) + wp.to_torch(ray_dirs_w)[:] = torch.tensor([[[0, 0, 1], [0, 0, 1]]], device=device) + wp.to_torch(ray_hits_w).fill_(float("inf")) + + wp.launch( + _raycast_mesh_masked_kernel, + dim=(n_envs, n_rays), + inputs=[mesh_id, env_mask, ray_starts_w, ray_dirs_w, float(1e6), 0, 0, ray_hits_w, ray_dist_w, ray_normal_w], + device=device, + ) + + torch.testing.assert_close(wp.to_torch(ray_hits_w), expected_hits) + assert torch.all(wp.to_torch(ray_dist_w) == _SENTINEL), "Distance buffer must not be written when return_distance=0" + assert torch.all(wp.to_torch(ray_normal_w) == _SENTINEL), "Normal buffer must not be written when return_normal=0" + + +def test_raycast_mesh_masked_kernel_with_distance(raycast_setup): + """return_distance=1: distances are written in addition to hits.""" + device = raycast_setup["device"] + mesh_id = raycast_setup["single_mesh_id"] + + n_envs, n_rays = 1, 2 + ray_starts_w, ray_dirs_w, ray_hits_w, ray_dist_w, ray_normal_w = _make_masked_buffers(device, n_envs, n_rays) + env_mask = wp.array([True], dtype=wp.bool, device=device) + + wp.to_torch(ray_starts_w)[:] = torch.tensor([[[0, -0.35, -5], [0.25, 0.35, -5]]], device=device) + wp.to_torch(ray_dirs_w)[:] = torch.tensor([[[0, 0, 1], [0, 0, 1]]], device=device) + wp.to_torch(ray_hits_w).fill_(float("inf")) + + wp.launch( + _raycast_mesh_masked_kernel, + dim=(n_envs, n_rays), + inputs=[mesh_id, env_mask, ray_starts_w, ray_dirs_w, float(1e6), 1, 0, ray_hits_w, ray_dist_w, ray_normal_w], + device=device, + ) + + # Cube bottom at z=-0.5, rays start at z=-5 going +z, distance = 4.5 + torch.testing.assert_close(wp.to_torch(ray_dist_w), torch.tensor([[4.5, 4.5]], device=device)) + assert torch.all(wp.to_torch(ray_normal_w) == _SENTINEL), "Normal buffer must not be written when return_normal=0" + + +def test_raycast_mesh_masked_kernel_with_normal(raycast_setup): + """return_distance=1, return_normal=1: both distances and surface normals are written.""" + device = raycast_setup["device"] + mesh_id = raycast_setup["single_mesh_id"] + + n_envs, n_rays = 1, 2 + ray_starts_w, ray_dirs_w, ray_hits_w, ray_dist_w, ray_normal_w = _make_masked_buffers(device, n_envs, n_rays) + env_mask = wp.array([True], dtype=wp.bool, device=device) + + wp.to_torch(ray_starts_w)[:] = torch.tensor([[[0, -0.35, -5], [0.25, 0.35, -5]]], device=device) + wp.to_torch(ray_dirs_w)[:] = torch.tensor([[[0, 0, 1], [0, 0, 1]]], device=device) + wp.to_torch(ray_hits_w).fill_(float("inf")) + + wp.launch( + _raycast_mesh_masked_kernel, + dim=(n_envs, n_rays), + inputs=[mesh_id, env_mask, ray_starts_w, ray_dirs_w, float(1e6), 1, 1, ray_hits_w, ray_dist_w, ray_normal_w], + device=device, + ) + + # Cube bottom at z=-0.5, rays start at z=-5, distance = 4.5 + torch.testing.assert_close(wp.to_torch(ray_dist_w), torch.tensor([[4.5, 4.5]], device=device)) + torch.testing.assert_close( + wp.to_torch(ray_normal_w), + torch.tensor([[[0, 0, -1], [0, 0, -1]]], device=device, dtype=torch.float32), + ) + + +def test_raycast_mesh_masked_kernel_env_mask(raycast_setup): + """Masked-out environments must not be written.""" + device = raycast_setup["device"] + mesh_id = raycast_setup["single_mesh_id"] + + n_envs, n_rays = 2, 2 + ray_starts_w, ray_dirs_w, ray_hits_w, ray_dist_w, ray_normal_w = _make_masked_buffers(device, n_envs, n_rays) + env_mask = wp.array([True, False], dtype=wp.bool, device=device) + + starts = torch.tensor([[[0, -0.35, -5], [0.25, 0.35, -5]], [[0, -0.35, -5], [0.25, 0.35, -5]]], device=device) + dirs = torch.tensor([[[0, 0, 1], [0, 0, 1]], [[0, 0, 1], [0, 0, 1]]], device=device) + wp.to_torch(ray_starts_w)[:] = starts + wp.to_torch(ray_dirs_w)[:] = dirs + wp.to_torch(ray_hits_w).fill_(float("inf")) + + wp.launch( + _raycast_mesh_masked_kernel, + dim=(n_envs, n_rays), + inputs=[mesh_id, env_mask, ray_starts_w, ray_dirs_w, float(1e6), 1, 0, ray_hits_w, ray_dist_w, ray_normal_w], + device=device, + ) + + hits = wp.to_torch(ray_hits_w) + dist = wp.to_torch(ray_dist_w) + + assert not torch.isinf(hits[0]).any(), "Active env 0 should have valid hits" + torch.testing.assert_close(dist[0], torch.tensor([4.5, 4.5], device=device)) + assert torch.isinf(hits[1]).all(), "Masked env 1 hits must remain inf" + assert torch.all(dist[1] == _SENTINEL), "Masked env 1 distances must remain at sentinel" + assert torch.all(wp.to_torch(ray_normal_w) == _SENTINEL), "Normal buffer must not be written when return_normal=0" + + +# --------------------------------------------------------------------------- +# Test quat_yaw_only correctness (regression for atan2-based fix) +# --------------------------------------------------------------------------- + + +@wp.kernel(enable_backward=False) +def _call_quat_yaw_only(q_in: wp.array(dtype=wp.quatf), q_out: wp.array(dtype=wp.quatf)): + i = wp.tid() + q_out[i] = _quat_yaw_only_func(q_in[i]) + + +def test_quat_yaw_only_pure_yaw(): + """Pure yaw: quat_yaw_only should match the yaw_quat() reference for all yaw angles.""" + device = "cuda" if torch.cuda.is_available() else "cpu" + yaw_angles = torch.tensor([0.0, 0.5, 1.2, -0.8, np.pi], device=device) + + for yaw in yaw_angles: + q_torch = quat_from_euler_xyz( + torch.tensor([0.0], device=device), + torch.tensor([0.0], device=device), + yaw.unsqueeze(0), + ) # shape (1, 4), xyzw + + expected = yaw_quat(q_torch) # shape (1, 4) + + q_in = wp.from_torch(q_torch.contiguous(), dtype=wp.quatf) + q_out = wp.zeros(1, dtype=wp.quatf, device=device) + wp.launch(_call_quat_yaw_only, dim=1, inputs=[q_in, q_out], device=device) + result = wp.to_torch(q_out) # shape (1, 4) + + torch.testing.assert_close(result, expected, atol=1e-5, rtol=1e-5) + + +def test_quat_yaw_only_with_pitch_roll(): + """Non-zero pitch and roll: only the yaw component should be preserved. + + This is the regression test for the old bug where simply zeroing qx/qy and + renormalizing gave the wrong answer when pitch or roll was non-zero. + """ + device = "cuda" if torch.cuda.is_available() else "cpu" + + # Several combined pitch+roll+yaw orientations: (roll, pitch, yaw) + test_cases = [ + (0.3, 0.4, 1.2), + (0.5, 0.0, 0.7), + (-0.2, 0.6, -1.0), + (1.0, 1.0, 0.0), # heavy pitch+roll, zero yaw → result should be identity + ] + + for roll, pitch, yaw in test_cases: + q_torch = quat_from_euler_xyz( + torch.tensor([roll], device=device), + torch.tensor([pitch], device=device), + torch.tensor([yaw], device=device), + ) # shape (1, 4), xyzw + + expected = yaw_quat(q_torch) # shape (1, 4) + + q_in = wp.from_torch(q_torch.contiguous(), dtype=wp.quatf) + q_out = wp.zeros(1, dtype=wp.quatf, device=device) + wp.launch(_call_quat_yaw_only, dim=1, inputs=[q_in, q_out], device=device) + result = wp.to_torch(q_out) + + torch.testing.assert_close(result, expected, atol=1e-5, rtol=1e-5) diff --git a/source/isaaclab/test/sensors/test_ray_caster_camera.py b/source/isaaclab/test/sensors/test_ray_caster_camera.py index c81ac9b2d74e..cc10b092a806 100644 --- a/source/isaaclab/test/sensors/test_ray_caster_camera.py +++ b/source/isaaclab/test/sensors/test_ray_caster_camera.py @@ -962,3 +962,142 @@ def test_sensor_print(setup_sim): sim.reset() # print info print(sensor) + + +@pytest.mark.isaacsim_ci +def test_depth_clipping_d2ip_and_d2c_are_independent(setup_sim): + """Clipping distance_to_image_plane must not corrupt distance_to_camera and vice versa. + + Both are derived from the same raw ray_distance buffer. If that buffer is modified + in-place by one clipping pass it would corrupt the other. This test verifies that + requesting both data types simultaneously gives results consistent with requesting + each one alone. + """ + sim, camera_cfg, dt = setup_sim + + base_cfg = RayCasterCameraCfg( + prim_path="/World/Camera", + mesh_prim_paths=["/World/defaultGroundPlane"], + offset=RayCasterCameraCfg.OffsetCfg(pos=(2.5, 2.5, 6.0), rot=(0.0, 0.1305, 0.0, 0.9914449), convention="world"), + pattern_cfg=patterns.PinholeCameraPatternCfg.from_intrinsic_matrix( + focal_length=38.0, + intrinsic_matrix=[380.08, 0.0, 467.79, 0.0, 380.08, 262.05, 0.0, 0.0, 1.0], + height=540, + width=960, + ), + max_distance=5.0, + data_types=["distance_to_image_plane", "distance_to_camera"], + depth_clipping_behavior="max", + update_period=0, + ) + + # Camera requesting both data types simultaneously + sim_utils.create_prim("/World/CameraJoint", "Xform") + cfg_joint = copy.deepcopy(base_cfg) + cfg_joint.prim_path = "/World/CameraJoint" + cam_joint = RayCasterCamera(cfg_joint) + + # Camera requesting only d2ip + sim_utils.create_prim("/World/CameraD2IP", "Xform") + cfg_d2ip = copy.deepcopy(base_cfg) + cfg_d2ip.prim_path = "/World/CameraD2IP" + cfg_d2ip.data_types = ["distance_to_image_plane"] + cam_d2ip = RayCasterCamera(cfg_d2ip) + + # Camera requesting only d2c + sim_utils.create_prim("/World/CameraD2C", "Xform") + cfg_d2c = copy.deepcopy(base_cfg) + cfg_d2c.prim_path = "/World/CameraD2C" + cfg_d2c.data_types = ["distance_to_camera"] + cam_d2c = RayCasterCamera(cfg_d2c) + + sim.reset() + + cam_joint.update(dt) + cam_d2ip.update(dt) + cam_d2c.update(dt) + + d2ip_joint = cam_joint.data.output["distance_to_image_plane"] + d2c_joint = cam_joint.data.output["distance_to_camera"] + d2ip_solo = cam_d2ip.data.output["distance_to_image_plane"] + d2c_solo = cam_d2c.data.output["distance_to_camera"] + + # Joint camera must match solo cameras (clipping one must not affect the other) + torch.testing.assert_close(d2ip_joint, d2ip_solo, atol=1e-5, rtol=1e-5) + torch.testing.assert_close(d2c_joint, d2c_solo, atol=1e-5, rtol=1e-5) + + # Both should be clipped to max_distance (camera is 6 m above ground, max_distance=5 m) + assert d2ip_joint.max().item() <= base_cfg.max_distance + 1e-4 + assert d2c_joint.max().item() <= base_cfg.max_distance + 1e-4 + + +@pytest.mark.isaacsim_ci +def test_frame_counter_increments_per_update(setup_sim): + """frame counter must increment by exactly 1 per update() call and reset to 0 on reset().""" + sim, camera_cfg, dt = setup_sim + camera = RayCasterCamera(cfg=camera_cfg) + sim.reset() + + assert torch.all(camera.frame == 0), "Frame must start at 0" + + n_steps = 7 + for step in range(1, n_steps + 1): + sim.step() + camera.update(dt, force_recompute=True) + assert camera.frame[0].item() == step, f"Frame must be {step} after {step} update(s)" + + # Partial reset: only env 0 (single-env camera, but API accepts env_ids) + camera.reset(env_ids=[0]) + assert camera.frame[0].item() == 0, "Frame must be 0 after reset(env_ids=[0])" + + # Full reset + for _ in range(3): + sim.step() + camera.update(dt, force_recompute=True) + camera.reset() + assert torch.all(camera.frame == 0), "Frame must be 0 after full reset()" + + +@pytest.mark.isaacsim_ci +def test_set_intrinsic_matrices_updates_output(setup_sim): + """Depth output must change when intrinsics are updated via set_intrinsic_matrices(). + + This tests that the warp view refresh in set_intrinsic_matrices() actually takes + effect: stale warp views would cause subsequent images to use the old ray pattern. + """ + sim, camera_cfg, dt = setup_sim + + # Place camera looking straight down at the ground + camera_cfg = copy.deepcopy(camera_cfg) + camera_cfg.offset = RayCasterCameraCfg.OffsetCfg(pos=(0.0, 0.0, 5.0), rot=(0.0, 0.0, 0.0, 1.0), convention="world") + camera_cfg.data_types = ["distance_to_camera"] + camera = RayCasterCamera(cfg=camera_cfg) + sim.reset() + + # Capture output with default focal length (24 mm → 20.955 mm aperture) + for _ in range(3): + sim.step() + camera.update(dt) + output_before = camera.data.output["distance_to_camera"].clone() + + # Change to a very different focal length (longer → tighter FOV → depth values differ at edges) + new_matrix = torch.tensor( + [[200.0, 0.0, 320.0], [0.0, 200.0, 240.0], [0.0, 0.0, 1.0]], + device=camera.device, + ).unsqueeze(0) + camera.set_intrinsic_matrices(new_matrix, focal_length=1.0) + + for _ in range(3): + sim.step() + camera.update(dt) + output_after = camera.data.output["distance_to_camera"].clone() + + # Outputs must differ after intrinsics change (different ray angles → different depths) + assert not torch.allclose(output_before, output_after, atol=1e-3), ( + "Depth output must change when intrinsic matrix is updated; unchanged output indicates stale warp ray buffers." + ) + # With depth_clipping_behavior="none" (default), missed rays produce inf — that is valid. + # No NaN values must appear; where rays hit, depth must be positive. + assert not torch.any(torch.isnan(output_after)), "Expected no NaN values in depth output after intrinsics update" + if torch.any(torch.isfinite(output_after)): + assert output_after[torch.isfinite(output_after)].min() > 0 diff --git a/source/isaaclab/test/sensors/test_ray_caster_integration.py b/source/isaaclab/test/sensors/test_ray_caster_integration.py new file mode 100644 index 000000000000..62b10a679661 --- /dev/null +++ b/source/isaaclab/test/sensors/test_ray_caster_integration.py @@ -0,0 +1,439 @@ +# 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 + +# pyright: reportPrivateUsage=none + +"""Integration tests for ray caster sensor view paths, env_mask, and intrinsics. + +These tests require Isaac Sim (AppLauncher). They cover the integration-level +items from ``TODO_ray_caster_kernel_tests.md``: + +- ``_get_view_transforms_wp`` ArticulationView and RigidBodyView paths +- ``MultiMeshRayCaster`` env_mask behavior +- ``MultiMeshRayCasterCamera.set_intrinsic_matrices`` propagation +- ``_update_mesh_transforms`` non-identity orientation offset (known bug, xfail) +- Depth clipping ordering for ``MultiMeshRayCasterCamera`` +""" + +from isaaclab.app import AppLauncher + +simulation_app = AppLauncher(headless=True, enable_cameras=True).app + +import copy + +import numpy as np +import pytest +import torch +import warp as wp + +from pxr import UsdGeom, UsdPhysics + +import isaaclab.sim as sim_utils +from isaaclab.sensors.ray_caster import ( + MultiMeshRayCaster, + MultiMeshRayCasterCamera, + MultiMeshRayCasterCameraCfg, + MultiMeshRayCasterCfg, + RayCaster, + RayCasterCfg, + patterns, +) +from isaaclab.terrains.trimesh.utils import make_plane +from isaaclab.terrains.utils import create_prim_from_mesh + +_GROUND_PATH = "/World/Ground" +_DT = 0.01 + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _make_sim_and_ground(): + """Create a blank stage with a flat ground plane at z=0.""" + sim_utils.create_new_stage() + sim = sim_utils.SimulationContext(sim_utils.SimulationCfg(dt=_DT)) + mesh = make_plane(size=(100, 100), height=0.0, center_zero=True) + create_prim_from_mesh(_GROUND_PATH, mesh) + sim_utils.update_stage() + return sim + + +def _single_downward_ray_cfg(prim_path: str) -> RayCasterCfg: + """RayCasterCfg with a single downward ray, no offset, world alignment.""" + return RayCasterCfg( + prim_path=prim_path, + mesh_prim_paths=[_GROUND_PATH], + update_period=0, + offset=RayCasterCfg.OffsetCfg(pos=(0.0, 0.0, 0.0), rot=(0.0, 0.0, 0.0, 1.0)), + debug_vis=False, + pattern_cfg=patterns.GridPatternCfg(resolution=1.0, size=(0.0, 0.0), direction=(0.0, 0.0, -1.0)), + ray_alignment="world", + ) + + +@pytest.fixture +def sim_ground(): + sim = _make_sim_and_ground() + yield sim + sim.stop() + sim.clear_instance() + + +# --------------------------------------------------------------------------- +# _get_view_transforms_wp: ArticulationView path +# --------------------------------------------------------------------------- + + +@pytest.mark.isaacsim_ci +def test_articulation_view_path(sim_ground): + """Mount a ray caster on a prim with ArticulationRootAPI. + + Verifies that sensor pos_w matches the prim's initial position and that + the downward ray hits the ground plane. This exercises the + ``ArticulationView.get_root_transforms()`` quaternion-convention path in + :meth:`_get_view_transforms_wp`. + """ + sim = sim_ground + expected_pos = (3.0, 4.0, 5.0) + + prim_path = "/World/ArticulatedBody" + sim_utils.create_prim(prim_path, "Xform", translation=expected_pos) + stage = sim_utils.get_current_stage() + prim = stage.GetPrimAtPath(prim_path) + UsdPhysics.RigidBodyAPI.Apply(prim) + UsdPhysics.ArticulationRootAPI.Apply(prim) + # Mass is needed for physics; collision is needed for PhysX to track the body. + mass_api = UsdPhysics.MassAPI.Apply(prim) + mass_api.CreateMassAttr().Set(1.0) + # Create a small collision cube so PhysX treats this as a real body. + cube_path = f"{prim_path}/CollisionCube" + cube_geom = UsdGeom.Cube.Define(stage, cube_path) + cube_geom.CreateSizeAttr().Set(0.1) + UsdPhysics.CollisionAPI.Apply(stage.GetPrimAtPath(cube_path)) + sim_utils.update_stage() + + sensor = RayCaster(_single_downward_ray_cfg(prim_path)) + sim.reset() + sensor.update(_DT) + + pos_w = wp.to_torch(sensor.data.pos_w)[0].cpu().numpy() + np.testing.assert_allclose( + pos_w, + expected_pos, + atol=0.15, + err_msg="ArticulationView: sensor pos_w must match initial prim position", + ) + + hits = wp.to_torch(sensor.data.ray_hits_w)[0, 0].cpu().numpy() + assert abs(hits[2]) < 0.5, f"ArticulationView: downward ray should hit near z=0, got z={hits[2]}" + + +# --------------------------------------------------------------------------- +# _get_view_transforms_wp: RigidBodyView path +# --------------------------------------------------------------------------- + + +@pytest.mark.isaacsim_ci +def test_rigid_body_view_path(sim_ground): + """Mount a ray caster on a prim with RigidBodyAPI (no ArticulationRootAPI). + + Exercises the ``RigidBodyView.get_transforms()`` path in + :meth:`_get_view_transforms_wp`. + """ + sim = sim_ground + expected_pos = (1.0, 2.0, 6.0) + + prim_path = "/World/RigidBody" + sim_utils.create_prim(prim_path, "Xform", translation=expected_pos) + stage = sim_utils.get_current_stage() + prim = stage.GetPrimAtPath(prim_path) + UsdPhysics.RigidBodyAPI.Apply(prim) + mass_api = UsdPhysics.MassAPI.Apply(prim) + mass_api.CreateMassAttr().Set(1.0) + cube_path = f"{prim_path}/CollisionCube" + cube_geom = UsdGeom.Cube.Define(stage, cube_path) + cube_geom.CreateSizeAttr().Set(0.1) + UsdPhysics.CollisionAPI.Apply(stage.GetPrimAtPath(cube_path)) + sim_utils.update_stage() + + sensor = RayCaster(_single_downward_ray_cfg(prim_path)) + sim.reset() + sensor.update(_DT) + + pos_w = wp.to_torch(sensor.data.pos_w)[0].cpu().numpy() + np.testing.assert_allclose( + pos_w, + expected_pos, + atol=0.15, + err_msg="RigidBodyView: sensor pos_w must match initial prim position", + ) + + hits = wp.to_torch(sensor.data.ray_hits_w)[0, 0].cpu().numpy() + assert abs(hits[2]) < 0.5, f"RigidBodyView: downward ray should hit near z=0, got z={hits[2]}" + + +# --------------------------------------------------------------------------- +# MultiMeshRayCasterCamera.set_intrinsic_matrices +# --------------------------------------------------------------------------- + + +@pytest.fixture +def sim_ground_camera(): + """Fixture providing sim + a base MultiMeshRayCasterCameraCfg.""" + sim = _make_sim_and_ground() + + camera_cfg = MultiMeshRayCasterCameraCfg( + prim_path="/World/Camera", + mesh_prim_paths=[_GROUND_PATH], + update_period=0, + offset=MultiMeshRayCasterCameraCfg.OffsetCfg(pos=(0.0, 0.0, 5.0), rot=(0.0, 0.0, 0.0, 1.0), convention="world"), + debug_vis=False, + pattern_cfg=patterns.PinholeCameraPatternCfg( + focal_length=24.0, + horizontal_aperture=20.955, + height=480, + width=640, + ), + data_types=["distance_to_camera"], + ) + + sim_utils.create_prim("/World/Camera", "Xform") + + yield sim, camera_cfg + + sim.stop() + sim.clear_instance() + + +@pytest.mark.isaacsim_ci +def test_multi_mesh_camera_set_intrinsic_matrices(sim_ground_camera): + """Depth output must change when intrinsics are updated on MultiMeshRayCasterCamera. + + The multi-mesh variant overrides ``_initialize_rays_impl`` without calling + ``super()``, so the warp view refresh path may differ from RayCasterCamera. + This test verifies that ``set_intrinsic_matrices`` actually takes effect. + """ + sim, camera_cfg = sim_ground_camera + + camera = MultiMeshRayCasterCamera(cfg=camera_cfg) + sim.reset() + + # Capture output with default intrinsics + for _ in range(3): + sim.step() + camera.update(_DT) + output_before = camera.data.output["distance_to_camera"].clone() + + # Change to a very different intrinsic matrix (different FOV) + new_matrix = torch.tensor( + [[200.0, 0.0, 320.0], [0.0, 200.0, 240.0], [0.0, 0.0, 1.0]], + device=camera.device, + ).unsqueeze(0) + camera.set_intrinsic_matrices(new_matrix, focal_length=1.0) + + for _ in range(3): + sim.step() + camera.update(_DT) + output_after = camera.data.output["distance_to_camera"].clone() + + assert not torch.allclose(output_before, output_after, atol=1e-3), ( + "MultiMeshRayCasterCamera: depth output must change after set_intrinsic_matrices; " + "unchanged output indicates stale warp ray buffers." + ) + assert not torch.any(torch.isnan(output_after)), "No NaN values expected after intrinsics update" + + +# --------------------------------------------------------------------------- +# Depth clipping ordering for MultiMeshRayCasterCamera +# --------------------------------------------------------------------------- + + +@pytest.mark.isaacsim_ci +def test_multi_mesh_camera_d2ip_and_d2c_independent(sim_ground_camera): + """Requesting both d2ip and d2c simultaneously must produce correct independent results. + + The ``distance_to_image_plane`` computation reads ``_ray_distance`` before + ``distance_to_camera`` clips it in-place. This test verifies the two data + types do not interfere with each other. + """ + sim, base_cfg = sim_ground_camera + + joint_cfg = copy.deepcopy(base_cfg) + joint_cfg.prim_path = "/World/CameraJoint" + joint_cfg.data_types = ["distance_to_image_plane", "distance_to_camera"] + joint_cfg.max_distance = 4.5 # camera is 5 m up, so some rays should be clipped + joint_cfg.depth_clipping_behavior = "max" + sim_utils.create_prim("/World/CameraJoint", "Xform") + cam_joint = MultiMeshRayCasterCamera(joint_cfg) + + d2ip_cfg = copy.deepcopy(base_cfg) + d2ip_cfg.prim_path = "/World/CameraD2IP" + d2ip_cfg.data_types = ["distance_to_image_plane"] + d2ip_cfg.max_distance = 4.5 + d2ip_cfg.depth_clipping_behavior = "max" + sim_utils.create_prim("/World/CameraD2IP", "Xform") + cam_d2ip = MultiMeshRayCasterCamera(d2ip_cfg) + + d2c_cfg = copy.deepcopy(base_cfg) + d2c_cfg.prim_path = "/World/CameraD2C" + d2c_cfg.data_types = ["distance_to_camera"] + d2c_cfg.max_distance = 4.5 + d2c_cfg.depth_clipping_behavior = "max" + sim_utils.create_prim("/World/CameraD2C", "Xform") + cam_d2c = MultiMeshRayCasterCamera(d2c_cfg) + + sim.reset() + + cam_joint.update(_DT) + cam_d2ip.update(_DT) + cam_d2c.update(_DT) + + d2ip_joint = cam_joint.data.output["distance_to_image_plane"] + d2c_joint = cam_joint.data.output["distance_to_camera"] + d2ip_solo = cam_d2ip.data.output["distance_to_image_plane"] + d2c_solo = cam_d2c.data.output["distance_to_camera"] + + # Joint camera must match solo cameras (clipping one must not corrupt the other) + torch.testing.assert_close(d2ip_joint, d2ip_solo, atol=1e-5, rtol=1e-5) + torch.testing.assert_close(d2c_joint, d2c_solo, atol=1e-5, rtol=1e-5) + + +# --------------------------------------------------------------------------- +# MultiMeshRayCaster env_mask behavior +# --------------------------------------------------------------------------- + + +@pytest.mark.isaacsim_ci +def test_multi_mesh_env_mask_preserves_masked_buffers(sim_ground): + """Masked environments must retain their pre-update buffer values. + + Creates a single-env MultiMeshRayCaster, captures output after one update, + then calls ``_update_buffers_impl`` with the environment masked out and + verifies the output buffers are unchanged. + """ + sim = sim_ground + + prim_path = "/World/Sensor" + sim_utils.create_prim(prim_path, "Xform", translation=(0.0, 0.0, 3.0)) + + cfg = MultiMeshRayCasterCfg( + prim_path=prim_path, + mesh_prim_paths=[_GROUND_PATH], + update_period=0, + offset=MultiMeshRayCasterCfg.OffsetCfg(pos=(0.0, 0.0, 0.0), rot=(0.0, 0.0, 0.0, 1.0)), + debug_vis=False, + pattern_cfg=patterns.GridPatternCfg(resolution=1.0, size=(0.0, 0.0), direction=(0.0, 0.0, -1.0)), + ray_alignment="world", + ) + sensor = MultiMeshRayCaster(cfg) + sim.reset() + + # First update: populate buffers with real values + sensor.update(_DT) + hits_before = wp.to_torch(sensor.data.ray_hits_w).clone() + + # Second update with env masked out: buffers must not change + mask_all_false = wp.array([False], dtype=wp.bool, device=sensor.device) + sensor._update_buffers_impl(mask_all_false) + + hits_after = wp.to_torch(sensor.data.ray_hits_w) + torch.testing.assert_close( + hits_after, + hits_before, + atol=0.0, + rtol=0.0, + msg="Masked env: ray_hits_w must be unchanged after update with env masked out", + ) + + +# --------------------------------------------------------------------------- +# _update_mesh_transforms: non-identity orientation offset +# --------------------------------------------------------------------------- + + +@pytest.mark.isaacsim_ci +def test_update_mesh_transforms_non_identity_offset(sim_ground): + """Tracked mesh position must account for body orientation when applying offset. + + Setup: a kinematic rigid body at (0, 0, 2) rotated 90 deg around Z, with a + child mesh offset by (1, 0, 0) in the body's local frame. + + Correct world position of mesh = body_pos + rotate(body_ori, local_offset) + = (0, 0, 2) + rotate(90degZ, (1, 0, 0)) + = (0, 0, 2) + (0, 1, 0) + = (0, 1, 2) + + Naive subtraction (the old bug) would give: body_pos - offset = (-1, 0, 2). + """ + sim = sim_ground + + from isaaclab.utils.math import quat_from_euler_xyz + + # 90 deg yaw quaternion in xyzw + yaw90 = quat_from_euler_xyz(torch.tensor([0.0]), torch.tensor([0.0]), torch.tensor([torch.pi / 2])) + yaw90_xyzw = tuple(yaw90[0].tolist()) + + # Create a kinematic rigid body at (0, 0, 2) rotated 90 deg around Z + body_path = "/World/DynamicBody" + sim_utils.create_prim(body_path, "Xform", translation=(0.0, 0.0, 2.0), orientation=yaw90_xyzw) + stage = sim_utils.get_current_stage() + body_prim = stage.GetPrimAtPath(body_path) + UsdPhysics.RigidBodyAPI.Apply(body_prim) + mass_api = UsdPhysics.MassAPI.Apply(body_prim) + mass_api.CreateMassAttr().Set(1.0) + body_prim.GetAttribute("physics:kinematicEnabled").Set(True) + + # Create a child Xform offset by (1, 0, 0) in the body's local frame, + # then place mesh geometry under it. The Xform translation is the offset + # that _obtain_trackable_prim_view / resolve_prim_pose will discover. + child_mesh_path = f"{body_path}/OffsetMesh" + sim_utils.create_prim(child_mesh_path, "Xform", translation=(1.0, 0.0, 0.0)) + mesh_data = make_plane(size=(2, 2), height=0.0, center_zero=True) + create_prim_from_mesh(f"{child_mesh_path}/Plane", mesh_data) + # Add collision so PhysX tracks the body + col_path = f"{body_path}/CollisionCube" + cube_geom = UsdGeom.Cube.Define(stage, col_path) + cube_geom.CreateSizeAttr().Set(0.1) + UsdPhysics.CollisionAPI.Apply(stage.GetPrimAtPath(col_path)) + sim_utils.update_stage() + + # Create a sensor prim to mount the MultiMeshRayCaster on + sensor_path = "/World/SensorMount" + sim_utils.create_prim(sensor_path, "Xform", translation=(0.0, 0.0, 5.0)) + + # Configure MultiMeshRayCaster to track the child mesh + cfg = MultiMeshRayCasterCfg( + prim_path=sensor_path, + mesh_prim_paths=[ + MultiMeshRayCasterCfg.RaycastTargetCfg( + prim_expr=child_mesh_path, + track_mesh_transforms=True, + ), + ], + update_period=0, + offset=MultiMeshRayCasterCfg.OffsetCfg(pos=(0.0, 0.0, 0.0), rot=(0.0, 0.0, 0.0, 1.0)), + debug_vis=False, + pattern_cfg=patterns.GridPatternCfg(resolution=1.0, size=(0.0, 0.0), direction=(0.0, 0.0, -1.0)), + ray_alignment="world", + ) + sensor = MultiMeshRayCaster(cfg) + sim.reset() + sensor.update(_DT) + + # Verify mesh position: body at (0,0,2) rotated 90deg Z, child offset (1,0,0) local + # Expected: (0, 0, 2) + rotate(90degZ, (1,0,0)) = (0, 0, 2) + (0, 1, 0) = (0, 1, 2) + mesh_pos = sensor._mesh_positions_w_torch.clone() + np.testing.assert_allclose( + mesh_pos[0, 0].cpu().numpy(), + [0.0, 1.0, 2.0], + atol=0.15, + err_msg=( + "Mesh position should be (0, 1, 2) via proper frame decomposition: " + "body_pos + rotate(body_ori, local_offset). " + "If this fails, the offset is not being rotated by the body orientation." + ), + ) diff --git a/source/isaaclab/test/sensors/test_ray_caster_kernels.py b/source/isaaclab/test/sensors/test_ray_caster_kernels.py new file mode 100644 index 000000000000..cc57e4f1eec5 --- /dev/null +++ b/source/isaaclab/test/sensors/test_ray_caster_kernels.py @@ -0,0 +1,577 @@ +# 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 + +"""Unit tests for ray caster kernels. + +Tests for kernels in ``sensors/ray_caster/kernels.py`` and +``utils/warp/kernels.py``. Exercised directly with hand-crafted warp arrays +and analytically computed expected outputs. No simulation, no stage, no +AppLauncher -- just warp and numpy on CPU (or CUDA when available). + +See ``test_update_ray_caster_kernel.py`` for tests of +:func:`update_ray_caster_kernel`. +""" + +from __future__ import annotations + +import importlib.util +import math +import os + +import numpy as np +import pytest +import warp as wp + +# --------------------------------------------------------------------------- +# Import kernel modules directly (avoids Isaac Sim / Omniverse dependencies) +# --------------------------------------------------------------------------- + +_SENSOR_KERNEL_PATH = os.path.join( + os.path.dirname(__file__), + os.pardir, + os.pardir, + "isaaclab", + "sensors", + "ray_caster", + "kernels.py", +) +_spec = importlib.util.spec_from_file_location("ray_caster_kernels", os.path.normpath(_SENSOR_KERNEL_PATH)) +_sensor_mod = importlib.util.module_from_spec(_spec) +_spec.loader.exec_module(_sensor_mod) + +_WARP_KERNEL_PATH = os.path.join( + os.path.dirname(__file__), + os.pardir, + os.pardir, + "isaaclab", + "utils", + "warp", + "kernels.py", +) +_warp_spec = importlib.util.spec_from_file_location("warp_kernels", os.path.normpath(_WARP_KERNEL_PATH)) +_warp_mod = importlib.util.module_from_spec(_warp_spec) +_warp_spec.loader.exec_module(_warp_mod) + +compute_distance_to_image_plane_masked_kernel = _sensor_mod.compute_distance_to_image_plane_masked_kernel +apply_depth_clipping_masked_kernel = _sensor_mod.apply_depth_clipping_masked_kernel +apply_z_drift_kernel = _sensor_mod.apply_z_drift_kernel +quat_yaw_only = _sensor_mod.quat_yaw_only + +raycast_dynamic_meshes_kernel = _warp_mod.raycast_dynamic_meshes_kernel + +# --------------------------------------------------------------------------- +# Constants & setup +# --------------------------------------------------------------------------- + +wp.init() +DEVICE = "cuda:0" if wp.is_cuda_available() else "cpu" +ATOL = 1e-5 + + +# --------------------------------------------------------------------------- +# Wrapper kernel for quat_yaw_only (@wp.func cannot be launched directly) +# --------------------------------------------------------------------------- + + +@wp.kernel(enable_backward=False) +def _quat_yaw_only_test_kernel( + q_in: wp.array(dtype=wp.quatf), + q_out: wp.array(dtype=wp.quatf), +): + tid = wp.tid() + q_out[tid] = quat_yaw_only(q_in[tid]) + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _euler_to_quat_xyzw(roll: float, pitch: float, yaw: float) -> tuple[float, float, float, float]: + """Euler angles (intrinsic XYZ) to quaternion in xyzw convention.""" + cr, sr = math.cos(roll / 2), math.sin(roll / 2) + cp, sp = math.cos(pitch / 2), math.sin(pitch / 2) + cy, sy = math.cos(yaw / 2), math.sin(yaw / 2) + qx = sr * cp * cy - cr * sp * sy + qy = cr * sp * cy + sr * cp * sy + qz = cr * cp * sy - sr * sp * cy + qw = cr * cp * cy + sr * sp * sy + return (qx, qy, qz, qw) + + +def _make_flat_mesh(size: float = 4.0) -> wp.Mesh: + """Create a flat square mesh in the XY plane at z=0, centered at origin.""" + half = size / 2.0 + vertices = np.array( + [[-half, -half, 0.0], [half, -half, 0.0], [half, half, 0.0], [-half, half, 0.0]], + dtype=np.float32, + ) + indices = np.array([0, 1, 2, 0, 2, 3], dtype=np.int32) + return wp.Mesh( + points=wp.array(vertices, dtype=wp.vec3, device=DEVICE), + indices=wp.array(indices, dtype=wp.int32, device=DEVICE), + ) + + +def _to_numpy(a: wp.array) -> np.ndarray: + """Convert a warp array to numpy, handling GPU arrays transparently.""" + return a.numpy() + + +# --------------------------------------------------------------------------- +# Tests: raycast_dynamic_meshes_kernel +# --------------------------------------------------------------------------- + + +class TestRaycastDynamicMeshesKernel: + """Tests for :func:`raycast_dynamic_meshes_kernel` from ``utils/warp/kernels.py``. + + Each test creates trivial warp meshes (flat quads) and verifies raycasting + results against analytical expectations. + """ + + IDENT_Q = [0.0, 0.0, 0.0, 1.0] + + @staticmethod + def _launch( + num_envs: int, + num_meshes: int, + num_rays: int, + env_mask: np.ndarray, + mesh_ids: np.ndarray, + ray_starts: np.ndarray, + ray_dirs: np.ndarray, + mesh_pos: np.ndarray, + mesh_rot: np.ndarray, + max_dist: float = 1e6, + sentinel: float | None = None, + ) -> dict[str, np.ndarray]: + """Build warp arrays, launch kernel, return outputs as numpy dicts.""" + env_mask_wp = wp.array(env_mask.astype(np.bool_), dtype=wp.bool, device=DEVICE) + mesh_wp = wp.array(mesh_ids, dtype=wp.uint64, device=DEVICE) + starts_wp = wp.array(ray_starts, dtype=wp.vec3f, device=DEVICE) + dirs_wp = wp.array(ray_dirs, dtype=wp.vec3f, device=DEVICE) + mpos_wp = wp.array(mesh_pos, dtype=wp.vec3f, device=DEVICE) + mrot_wp = wp.array(mesh_rot, dtype=wp.quatf, device=DEVICE) + + fill = sentinel if sentinel is not None else float("inf") + + hits_np = np.full((num_envs, num_rays, 3), fill, dtype=np.float32) + ray_hits = wp.array(hits_np, dtype=wp.vec3f, device=DEVICE) + + dist_np = np.full((num_envs, num_rays), fill, dtype=np.float32) + ray_distance = wp.array(dist_np, dtype=wp.float32, device=DEVICE) + + normal_np = np.full((num_envs, num_rays, 3), fill, dtype=np.float32) + ray_normal = wp.array(normal_np, dtype=wp.vec3f, device=DEVICE) + + face_np = np.full((num_envs, num_rays), -1, dtype=np.int32) + ray_face_id = wp.array(face_np, dtype=wp.int32, device=DEVICE) + + mesh_id_np = np.full((num_envs, num_rays), -1, dtype=np.int16) + ray_mesh_id = wp.array(mesh_id_np, dtype=wp.int16, device=DEVICE) + + wp.launch( + raycast_dynamic_meshes_kernel, + dim=(num_meshes, num_envs, num_rays), + inputs=[ + env_mask_wp, + mesh_wp, + starts_wp, + dirs_wp, + ray_hits, + ray_distance, + ray_normal, + ray_face_id, + ray_mesh_id, + mpos_wp, + mrot_wp, + max_dist, + 1, # return_normal + 1, # return_face_id + 1, # return_mesh_id + ], + device=DEVICE, + ) + wp.synchronize_device(DEVICE) + + return { + "hits": _to_numpy(ray_hits), + "distance": _to_numpy(ray_distance), + "normal": _to_numpy(ray_normal), + "face_id": _to_numpy(ray_face_id), + "mesh_id": _to_numpy(ray_mesh_id), + } + + def test_env_mask_skipping(self): + """Env 0 masked out -- verify output buffers retain sentinel values.""" + mesh = _make_flat_mesh() + iq = self.IDENT_Q + out = self._launch( + num_envs=2, + num_meshes=1, + num_rays=1, + env_mask=np.array([False, True]), + mesh_ids=np.array([[mesh.id], [mesh.id]], dtype=np.uint64), + ray_starts=np.array([[[0, 0, 10]], [[0, 0, 10]]], dtype=np.float32), + ray_dirs=np.array([[[0, 0, -1]], [[0, 0, -1]]], dtype=np.float32), + mesh_pos=np.array([[[0, 0, 2]], [[0, 0, 2]]], dtype=np.float32), + mesh_rot=np.array([[iq], [iq]], dtype=np.float32), + sentinel=999.0, + ) + + # Env 0 (masked): all outputs retain sentinel / initial fill + np.testing.assert_allclose(out["hits"][0, 0], [999, 999, 999], atol=ATOL) + assert out["distance"][0, 0] == pytest.approx(999.0, abs=ATOL) + np.testing.assert_allclose(out["normal"][0, 0], [999, 999, 999], atol=ATOL) + assert out["face_id"][0, 0] == -1 + assert out["mesh_id"][0, 0] == -1 + + # Env 1 (active): should have hit the mesh at z=2, distance 8 + np.testing.assert_allclose(out["hits"][1, 0], [0, 0, 2], atol=ATOL) + assert out["distance"][1, 0] == pytest.approx(8.0, abs=ATOL) + assert out["mesh_id"][1, 0] == 0 + + def test_closest_hit_overlapping_meshes(self): + """Two meshes at different distances -- closer hit wins. + + Mesh A at z=2 (farther), Mesh B at z=4 (closer to ray origin at z=10). + Ray from (0,0,10) going (0,0,-1). Expected: hit Mesh B at distance 6. + """ + mesh_a = _make_flat_mesh() + mesh_b = _make_flat_mesh() + iq = self.IDENT_Q + + out = self._launch( + num_envs=1, + num_meshes=2, + num_rays=1, + env_mask=np.array([True]), + mesh_ids=np.array([[mesh_a.id, mesh_b.id]], dtype=np.uint64), + ray_starts=np.array([[[0, 0, 10]]], dtype=np.float32), + ray_dirs=np.array([[[0, 0, -1]]], dtype=np.float32), + mesh_pos=np.array([[[0, 0, 2], [0, 0, 4]]], dtype=np.float32), + mesh_rot=np.array([[iq, iq]], dtype=np.float32), + ) + + np.testing.assert_allclose(out["hits"][0, 0], [0, 0, 4], atol=ATOL) + assert out["distance"][0, 0] == pytest.approx(6.0, abs=ATOL) + np.testing.assert_allclose(out["normal"][0, 0], [0, 0, 1], atol=ATOL) + assert out["mesh_id"][0, 0] == 1 # mesh_b is closer + + def test_mesh_transform_application(self): + """Mesh translated/rotated -- verify hits in correct world-space coordinates. + + Mesh: flat XY quad at z=0 (local), placed at world (5,0,0) with 90 deg + Y rotation. This turns it into a vertical plane at x=5. + Ray from (10,0,0) going (-1,0,0) should hit at (5,0,0), distance=5. + World-space normal: local (0,0,1) rotated by 90 deg Y = (1,0,0). + """ + mesh = _make_flat_mesh() + rot90y = [0.0, math.sin(math.pi / 4), 0.0, math.cos(math.pi / 4)] + + out = self._launch( + num_envs=1, + num_meshes=1, + num_rays=1, + env_mask=np.array([True]), + mesh_ids=np.array([[mesh.id]], dtype=np.uint64), + ray_starts=np.array([[[10, 0, 0]]], dtype=np.float32), + ray_dirs=np.array([[[-1, 0, 0]]], dtype=np.float32), + mesh_pos=np.array([[[5, 0, 0]]], dtype=np.float32), + mesh_rot=np.array([[rot90y]], dtype=np.float32), + ) + + np.testing.assert_allclose(out["hits"][0, 0], [5, 0, 0], atol=ATOL) + assert out["distance"][0, 0] == pytest.approx(5.0, abs=ATOL) + np.testing.assert_allclose(out["normal"][0, 0], [1, 0, 0], atol=ATOL) + + def test_equidistant_meshes(self): + """Two meshes at exact same distance -- hit position is always correct. + + Known limitation (warp#1058): when two meshes are equidistant, the + ``atomic_min`` + equality-check pattern is not fully thread-safe. + Normals, face_ids, and mesh_ids may come from either mesh. The hit + *position* is always correct because both threads compute the same + world-space point. + """ + mesh_a = _make_flat_mesh() + mesh_b = _make_flat_mesh() + iq = self.IDENT_Q + + out = self._launch( + num_envs=1, + num_meshes=2, + num_rays=1, + env_mask=np.array([True]), + mesh_ids=np.array([[mesh_a.id, mesh_b.id]], dtype=np.uint64), + ray_starts=np.array([[[0, 0, 10]]], dtype=np.float32), + ray_dirs=np.array([[[0, 0, -1]]], dtype=np.float32), + mesh_pos=np.array([[[0, 0, 3], [0, 0, 3]]], dtype=np.float32), + mesh_rot=np.array([[iq, iq]], dtype=np.float32), + ) + + # Position and distance are always correct, even under the race + np.testing.assert_allclose(out["hits"][0, 0], [0, 0, 3], atol=ATOL) + assert out["distance"][0, 0] == pytest.approx(7.0, abs=ATOL) + # mesh_id can be 0 or 1 -- both are valid under the race condition + assert out["mesh_id"][0, 0] in (0, 1) + + +# --------------------------------------------------------------------------- +# Tests: compute_distance_to_image_plane_masked_kernel +# --------------------------------------------------------------------------- + + +class TestComputeDistanceToImagePlaneMaskedKernel: + """Tests for :func:`compute_distance_to_image_plane_masked_kernel`.""" + + @staticmethod + def _launch( + quat_xyzw: list[float], + ray_distance: list[list[float]], + ray_dirs: list[list[list[float]]], + env_mask: list[bool] | None = None, + ) -> np.ndarray: + """Launch kernel and return distance_to_image_plane as numpy.""" + num_envs = len(ray_distance) + num_rays = len(ray_distance[0]) + if env_mask is None: + env_mask = [True] * num_envs + + mask_wp = wp.array(np.array(env_mask, dtype=np.bool_), dtype=wp.bool, device=DEVICE) + quat_np = np.array([quat_xyzw] * num_envs, dtype=np.float32) + quat_wp = wp.array(quat_np, dtype=wp.quatf, device=DEVICE) + ray_dist_wp = wp.array(np.array(ray_distance, dtype=np.float32), dtype=wp.float32, device=DEVICE) + dirs_wp = wp.array(np.array(ray_dirs, dtype=np.float32), dtype=wp.vec3f, device=DEVICE) + out_wp = wp.zeros((num_envs, num_rays), dtype=wp.float32, device=DEVICE) + + wp.launch( + compute_distance_to_image_plane_masked_kernel, + dim=(num_envs, num_rays), + inputs=[mask_wp, quat_wp, ray_dist_wp, dirs_wp], + outputs=[out_wp], + device=DEVICE, + ) + wp.synchronize_device(DEVICE) + return _to_numpy(out_wp) + + def test_known_camera_orientation(self): + """Identity camera, ray along +X at distance 5 -- d2ip equals 5.""" + result = self._launch( + quat_xyzw=[0, 0, 0, 1], + ray_distance=[[5.0]], + ray_dirs=[[[1, 0, 0]]], + ) + assert result[0, 0] == pytest.approx(5.0, abs=ATOL) + + def test_off_axis_camera(self): + """Camera pitched 45 deg around Y, ray going world -Z. + + Camera forward (+X_cam) in world = (cos45, 0, -sin45). + Displacement = 10 * (0, 0, -1) = (0, 0, -10). + Projection onto camera forward = dot((0,0,-10), (cos45,0,-sin45)) + = 10 * sin(45 deg). + """ + pitch45 = list(_euler_to_quat_xyzw(0, math.pi / 4, 0)) + result = self._launch( + quat_xyzw=pitch45, + ray_distance=[[10.0]], + ray_dirs=[[[0, 0, -1]]], + ) + expected = 10.0 * math.sin(math.pi / 4) + assert result[0, 0] == pytest.approx(expected, abs=ATOL) + + def test_inf_distance(self): + """Inf distance produces NaN through the projection (inf * 0 = NaN). + + When a ray misses, ray_distance is inf. Multiplying inf by zero-valued + ray-direction components yields NaN (IEEE 754), which propagates through + the quaternion rotation. The downstream + :func:`apply_depth_clipping_masked_kernel` handles NaN correctly via + ``wp.isnan()``, so the overall pipeline is sound. + """ + result = self._launch( + quat_xyzw=[0, 0, 0, 1], + ray_distance=[[float("inf")]], + ray_dirs=[[[1, 0, 0]]], + ) + assert np.isnan(result[0, 0]), f"Expected NaN from inf*0 contamination, got {result[0, 0]}" + + +# --------------------------------------------------------------------------- +# Tests: apply_depth_clipping_masked_kernel +# --------------------------------------------------------------------------- + + +class TestApplyDepthClippingMaskedKernel: + """Tests for :func:`apply_depth_clipping_masked_kernel`.""" + + @staticmethod + def _launch( + depth_values: list[list[float]], + max_dist: float, + fill_val: float, + env_mask: list[bool] | None = None, + ) -> np.ndarray: + """Launch kernel and return clipped depth as numpy.""" + num_envs = len(depth_values) + num_rays = len(depth_values[0]) + if env_mask is None: + env_mask = [True] * num_envs + + mask_wp = wp.array(np.array(env_mask, dtype=np.bool_), dtype=wp.bool, device=DEVICE) + depth_wp = wp.array(np.array(depth_values, dtype=np.float32), dtype=wp.float32, device=DEVICE) + + wp.launch( + apply_depth_clipping_masked_kernel, + dim=(num_envs, num_rays), + inputs=[mask_wp, max_dist, fill_val], + outputs=[depth_wp], + device=DEVICE, + ) + wp.synchronize_device(DEVICE) + return _to_numpy(depth_wp) + + def test_boundary_at_max_dist(self): + """Value at exactly max_dist is preserved (not clipped).""" + result = self._launch([[10.0]], max_dist=10.0, fill_val=0.0) + assert result[0, 0] == pytest.approx(10.0, abs=ATOL) + + def test_above_max_dist(self): + """Value above max_dist is replaced with fill_val.""" + result = self._launch([[10.001]], max_dist=10.0, fill_val=0.0) + assert result[0, 0] == pytest.approx(0.0, abs=ATOL) + + def test_nan_value(self): + """NaN value is replaced with fill_val.""" + result = self._launch([[float("nan")]], max_dist=10.0, fill_val=0.0) + assert result[0, 0] == pytest.approx(0.0, abs=ATOL) + + def test_inf_value(self): + """Inf is clipped (inf > max_dist is true).""" + result = self._launch([[float("inf")]], max_dist=10.0, fill_val=0.0) + assert result[0, 0] == pytest.approx(0.0, abs=ATOL) + + def test_negative_depth(self): + """Negative depth passes through unclipped (valid for distance-to-image-plane).""" + result = self._launch([[-3.5]], max_dist=10.0, fill_val=0.0) + assert result[0, 0] == pytest.approx(-3.5, abs=ATOL) + + def test_env_mask(self): + """Masked env retains original value -- clipping is not applied.""" + result = self._launch( + depth_values=[[15.0], [15.0]], + max_dist=10.0, + fill_val=0.0, + env_mask=[False, True], + ) + # Env 0 (masked): unchanged + assert result[0, 0] == pytest.approx(15.0, abs=ATOL) + # Env 1 (active): clipped + assert result[1, 0] == pytest.approx(0.0, abs=ATOL) + + def test_fill_val_zero_vs_max(self): + """fill_val=0.0 and fill_val=max_dist produce correct replacements.""" + max_dist = 10.0 + + result_zero = self._launch([[15.0]], max_dist=max_dist, fill_val=0.0) + assert result_zero[0, 0] == pytest.approx(0.0, abs=ATOL) + + result_max = self._launch([[15.0]], max_dist=max_dist, fill_val=max_dist) + assert result_max[0, 0] == pytest.approx(max_dist, abs=ATOL) + + +# --------------------------------------------------------------------------- +# Tests: apply_z_drift_kernel +# --------------------------------------------------------------------------- + + +class TestApplyZDriftKernel: + """Tests for :func:`apply_z_drift_kernel`.""" + + @staticmethod + def _launch( + hits: list[list[list[float]]], + drift: list[list[float]], + env_mask: list[bool] | None = None, + ) -> np.ndarray: + """Launch kernel and return modified ray_hits as numpy.""" + num_envs = len(hits) + num_rays = len(hits[0]) + if env_mask is None: + env_mask = [True] * num_envs + + mask_wp = wp.array(np.array(env_mask, dtype=np.bool_), dtype=wp.bool, device=DEVICE) + drift_wp = wp.array(np.array(drift, dtype=np.float32), dtype=wp.vec3f, device=DEVICE) + hits_wp = wp.array(np.array(hits, dtype=np.float32), dtype=wp.vec3f, device=DEVICE) + + wp.launch( + apply_z_drift_kernel, + dim=(num_envs, num_rays), + inputs=[mask_wp, drift_wp], + outputs=[hits_wp], + device=DEVICE, + ) + wp.synchronize_device(DEVICE) + return _to_numpy(hits_wp) + + def test_known_drift(self): + """ray_cast_drift = (0, 0, 1.5) shifts ray hit z by exactly 1.5.""" + result = self._launch( + hits=[[[3.0, 4.0, 5.0]]], + drift=[[0.0, 0.0, 1.5]], + ) + np.testing.assert_allclose(result[0, 0], [3.0, 4.0, 6.5], atol=ATOL) + + def test_only_z_component(self): + """Only z-component of drift is applied; x and y are unchanged.""" + result = self._launch( + hits=[[[3.0, 4.0, 5.0]]], + drift=[[0.5, 0.3, 1.0]], + ) + np.testing.assert_allclose(result[0, 0], [3.0, 4.0, 6.0], atol=ATOL) + + +# --------------------------------------------------------------------------- +# Tests: quat_yaw_only +# --------------------------------------------------------------------------- + + +class TestQuatYawOnly: + """Tests for :func:`quat_yaw_only` (a ``@wp.func`` tested via wrapper kernel).""" + + def test_gimbal_lock(self): + """At pitch = +/-pi/2, atan2 is near-degenerate but should produce a + finite, unit-norm, pure-yaw quaternion (only z and w components). + """ + q_down = _euler_to_quat_xyzw(0, math.pi / 2, 0) # pitch = +pi/2 + q_up = _euler_to_quat_xyzw(0, -math.pi / 2, 0) # pitch = -pi/2 + + q_in_np = np.array([list(q_down), list(q_up)], dtype=np.float32) + q_in = wp.array(q_in_np, dtype=wp.quatf, device=DEVICE) + q_out = wp.zeros(2, dtype=wp.quatf, device=DEVICE) + + wp.launch( + _quat_yaw_only_test_kernel, + dim=2, + inputs=[q_in], + outputs=[q_out], + device=DEVICE, + ) + wp.synchronize_device(DEVICE) + + result = _to_numpy(q_out) + + for i in range(2): + qx, qy, qz, qw = result[i] + # Must be finite (no NaN / inf) + assert np.isfinite(result[i]).all(), f"Non-finite output at index {i}: {result[i]}" + # Must be a pure-yaw quaternion: x ~ 0, y ~ 0 + assert abs(qx) < ATOL, f"x-component should be ~0 at gimbal lock, got {qx}" + assert abs(qy) < ATOL, f"y-component should be ~0 at gimbal lock, got {qy}" + # Must be unit-norm + norm = math.sqrt(float(qx) ** 2 + float(qy) ** 2 + float(qz) ** 2 + float(qw) ** 2) + assert norm == pytest.approx(1.0, abs=ATOL), f"Non-unit quaternion at index {i}: norm={norm}" diff --git a/source/isaaclab/test/sensors/test_ray_caster_sensor.py b/source/isaaclab/test/sensors/test_ray_caster_sensor.py new file mode 100644 index 000000000000..f1c90a986b5d --- /dev/null +++ b/source/isaaclab/test/sensors/test_ray_caster_sensor.py @@ -0,0 +1,272 @@ +# 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 + +# pyright: reportPrivateUsage=none + +"""Tests for RayCaster sensor behavior: alignment modes and reset.""" + +from isaaclab.app import AppLauncher + +simulation_app = AppLauncher(headless=True).app + +import numpy as np +import pytest +import torch +import warp as wp + +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 +from isaaclab.utils.math import quat_from_euler_xyz + +# ------------------------------------------------------------------- +# Helpers +# ------------------------------------------------------------------- + +_GROUND_PATH = "/World/Ground" + + +def _make_sim_and_ground(): + """Create a blank stage with a flat ground plane at z=0 and return the SimulationContext.""" + sim_utils.create_new_stage() + dt = 0.01 + sim_cfg = sim_utils.SimulationCfg(dt=dt) + sim = sim_utils.SimulationContext(sim_cfg) + mesh = make_plane(size=(100, 100), height=0.0, center_zero=True) + create_prim_from_mesh(_GROUND_PATH, mesh) + sim_utils.update_stage() + return sim + + +def _ray_caster_cfg(prim_path: str, alignment: str) -> RayCasterCfg: + """Single downward ray, no offset from prim.""" + return RayCasterCfg( + prim_path=prim_path, + mesh_prim_paths=[_GROUND_PATH], + update_period=0, + offset=RayCasterCfg.OffsetCfg(pos=(0.0, 0.0, 0.0), rot=(0.0, 0.0, 0.0, 1.0)), + debug_vis=False, + pattern_cfg=patterns.GridPatternCfg(resolution=1.0, size=(0.0, 0.0), direction=(0.0, 0.0, -1.0)), + ray_alignment=alignment, + ) + + +@pytest.fixture +def sim_ground(): + sim = _make_sim_and_ground() + yield sim + sim.stop() + sim.clear_instance() + + +# ------------------------------------------------------------------- +# Alignment mode tests +# ------------------------------------------------------------------- + + +@pytest.mark.isaacsim_ci +def test_world_alignment_ignores_sensor_pitch(sim_ground): + """In 'world' alignment, ray direction is always (0,0,-1) regardless of sensor pitch. + + Two sensors at the same location: one upright (identity), one pitched 30°. + World-mode sensors must produce the same hit position (straight below at z=0). + """ + sim = sim_ground + + # Upright sensor: identity orientation + sim_utils.create_prim("/World/SensorUpright", "Xform", translation=(0.0, 0.0, 2.0)) + # Pitched 30° sensor — orientation=(x,y,z,w) per Isaac Lab convention + pitch_quat = quat_from_euler_xyz( + torch.tensor([0.0]), torch.tensor([np.pi / 6]), torch.tensor([0.0]) + ) # shape (1, 4), xyzw + sim_utils.create_prim( + "/World/SensorPitched", + "Xform", + translation=(0.0, 0.0, 2.0), + orientation=tuple(pitch_quat[0].tolist()), # xyzw + ) + + sensor_upright = RayCaster(_ray_caster_cfg("/World/SensorUpright", "world")) + sensor_pitched = RayCaster(_ray_caster_cfg("/World/SensorPitched", "world")) + sim.reset() + + dt = 0.01 + sensor_upright.update(dt) + sensor_pitched.update(dt) + + # ray_hits_w is a wp.array(dtype=wp.vec3f); convert to torch for indexing + hits_upright = wp.to_torch(sensor_upright.data.ray_hits_w) # (1, 1, 3) + hits_pitched = wp.to_torch(sensor_pitched.data.ray_hits_w) + + # Both must hit z=0 (straight down, world frame direction) + assert abs(hits_upright[0, 0, 2].item()) < 0.02, ( + f"Upright world sensor must hit z≈0, got {hits_upright[0, 0, 2].item()}" + ) + assert abs(hits_pitched[0, 0, 2].item()) < 0.02, ( + f"Pitched world sensor must hit z≈0, got {hits_pitched[0, 0, 2].item()}" + ) + # Lateral positions must agree (same start at [0,0,2] + same direction [0,0,-1]) + torch.testing.assert_close(hits_upright, hits_pitched, atol=0.02, rtol=0) + + +@pytest.mark.isaacsim_ci +def test_base_alignment_rotates_ray_direction(sim_ground): + """In 'base' alignment, ray direction follows the full sensor orientation. + + A sensor pitched +30° around Y (quat_from_euler_xyz(pitch=pi/6)): + - Rotates (0,0,-1) to (-sin(30°), 0, -cos(30°)) = (-0.5, 0, -0.866) + - world mode → ray still goes straight down, hits x≈0, z≈0 + - base mode → ray tilts, hits at x ≈ -2*tan(30°) ≈ -1.155 + """ + sim = sim_ground + + pitch_quat = quat_from_euler_xyz( + torch.tensor([0.0]), torch.tensor([np.pi / 6]), torch.tensor([0.0]) + ) # shape (1, 4), xyzw + orientation = tuple(pitch_quat[0].tolist()) + + sim_utils.create_prim("/World/SensorWorld", "Xform", translation=(0.0, 0.0, 2.0), orientation=orientation) + sim_utils.create_prim("/World/SensorBase", "Xform", translation=(0.0, 0.0, 2.0), orientation=orientation) + + sensor_world = RayCaster(_ray_caster_cfg("/World/SensorWorld", "world")) + sensor_base = RayCaster(_ray_caster_cfg("/World/SensorBase", "base")) + sim.reset() + + dt = 0.01 + sensor_world.update(dt) + sensor_base.update(dt) + + hits_world = wp.to_torch(sensor_world.data.ray_hits_w) # (1, 1, 3) + hits_base = wp.to_torch(sensor_base.data.ray_hits_w) + + # World mode: ray still hits directly below (x≈0, y≈0, z≈0) + assert abs(hits_world[0, 0, 0].item()) < 0.05, f"World mode hit x must be near 0, got {hits_world[0, 0, 0].item()}" + assert abs(hits_world[0, 0, 2].item()) < 0.05, f"World mode must hit z≈0, got {hits_world[0, 0, 2].item()}" + + # Base mode: pitch +30° around Y rotates (0,0,-1) to (-0.5, 0, -0.866). + # From height 2, the ray hits x = -2 * tan(30°) ≈ -1.155. + expected_x = -2.0 * np.tan(np.pi / 6) + assert abs(hits_base[0, 0, 0].item() - expected_x) < 0.05, ( + f"Base mode hit x should be ≈{expected_x:.3f}, got {hits_base[0, 0, 0].item():.3f}" + ) + assert abs(hits_base[0, 0, 2].item()) < 0.05, f"Base mode must hit ground (z≈0), got {hits_base[0, 0, 2].item()}" + + +@pytest.mark.isaacsim_ci +def test_yaw_alignment_direction_unchanged(sim_ground): + """In 'yaw' alignment, ray directions stay world-down despite pitch+roll. + + Setup: sensor at (0,0,2), pitched 30° and yawed 45°; pattern has a single ray + at local offset (+1, 0, 0). + + - world mode: start = sensor_pos + (1,0,0) (no rotation applied to offset) + - yaw mode: start = sensor_pos + yaw_rot(45°) @ (1,0,0) = (cos45°, sin45°, 0) + + Both modes fire the ray straight down (direction unchanged), so both hit z=0. + The hit x-coordinate differs between modes, confirming the yaw-only rotation of + start positions is applied in 'yaw' mode but not in 'world' mode. + """ + sim = sim_ground + + combined_quat = quat_from_euler_xyz( + torch.tensor([0.0]), + torch.tensor([np.pi / 6]), # 30° pitch + torch.tensor([np.pi / 4]), # 45° yaw + ) # shape (1, 4), xyzw + orientation = tuple(combined_quat[0].tolist()) + + sim_utils.create_prim("/World/SensorWorldY", "Xform", translation=(0.0, 0.0, 2.0), orientation=orientation) + sim_utils.create_prim("/World/SensorYaw", "Xform", translation=(0.0, 0.0, 2.0), orientation=orientation) + + # Use a single ray at local offset (+1, 0, 0), still pointing down + def _cfg_with_offset(prim_path, alignment): + return RayCasterCfg( + prim_path=prim_path, + mesh_prim_paths=[_GROUND_PATH], + update_period=0, + offset=RayCasterCfg.OffsetCfg(pos=(1.0, 0.0, 0.0), rot=(0.0, 0.0, 0.0, 1.0)), + debug_vis=False, + pattern_cfg=patterns.GridPatternCfg(resolution=1.0, size=(0.0, 0.0), direction=(0.0, 0.0, -1.0)), + ray_alignment=alignment, + ) + + sensor_world = RayCaster(_cfg_with_offset("/World/SensorWorldY", "world")) + sensor_yaw = RayCaster(_cfg_with_offset("/World/SensorYaw", "yaw")) + sim.reset() + + dt = 0.01 + sensor_world.update(dt) + sensor_yaw.update(dt) + + hits_world = wp.to_torch(sensor_world.data.ray_hits_w) # (1, 1, 3) + hits_yaw = wp.to_torch(sensor_yaw.data.ray_hits_w) + + # Both modes must hit the ground (direction unchanged = straight down in both modes) + assert abs(hits_world[0, 0, 2].item()) < 0.05, "World mode must hit z≈0" + assert abs(hits_yaw[0, 0, 2].item()) < 0.05, "Yaw mode must hit z≈0 (direction straight down)" + + # world mode: offset (1,0,0) not rotated → ray starts at sensor_pos+(1,0,0) → hits x≈1 + assert abs(hits_world[0, 0, 0].item() - 1.0) < 0.05, ( + f"World mode: hit x should be ≈1.0 (unrotated offset), got {hits_world[0, 0, 0].item():.3f}" + ) + + # yaw mode: offset (1,0,0) rotated by 45° yaw → starts at sensor_pos+(cos45°, sin45°, 0) → hits x≈cos45° + expected_x_yaw = np.cos(np.pi / 4) # ≈ 0.707 + assert abs(hits_yaw[0, 0, 0].item() - expected_x_yaw) < 0.05, ( + f"Yaw mode: hit x should be ≈{expected_x_yaw:.3f} (yaw-rotated offset), got {hits_yaw[0, 0, 0].item():.3f}" + ) + # Confirm they differ — if they were the same, the test would not cover the yaw rotation + assert not torch.allclose(hits_world, hits_yaw, atol=0.1), ( + "Yaw and world modes must produce different hit positions for non-zero lateral offset" + ) + + +# ------------------------------------------------------------------- +# Reset / drift test +# ------------------------------------------------------------------- + + +@pytest.mark.isaacsim_ci +def test_ray_caster_reset_resamples_drift(sim_ground): + """reset() resamples drift values within the configured drift_range.""" + sim = sim_ground + + sim_utils.create_prim("/World/Sensor", "Xform", translation=(0.0, 0.0, 2.0)) + cfg = _ray_caster_cfg("/World/Sensor", "world") + cfg.drift_range = (0.01, 0.05) # force non-zero drift + sensor = RayCaster(cfg) + sim.reset() + # sim.reset() initializes the sensor with zero drift; call sensor.reset() to resample + # from the configured drift_range before we capture the baseline. + sensor.reset() + + dt = 0.01 + sensor.update(dt) + drift_before = sensor.drift.clone() # (1, 3) torch tensor + + lo, hi = cfg.drift_range + + # After sensor.reset(), drift should be within the configured range + assert drift_before.shape == (1, 3), f"Drift shape should be (1, 3), got {drift_before.shape}" + assert (drift_before >= lo - 1e-6).all() and (drift_before <= hi + 1e-6).all(), ( + f"Initial drift must be in [{lo}, {hi}], got [{drift_before.min():.4f}, {drift_before.max():.4f}]" + ) + + # reset() resamples drift; values should remain within the configured range + # Call reset() multiple times until we get a different sample (probability of same is near zero + # for continuous uniform distribution, but we retry to avoid flakiness). + for _ in range(5): + sensor.reset() + drift_after = sensor.drift.clone() + if not torch.allclose(drift_after, drift_before): + break + assert drift_after.shape == drift_before.shape, "Drift shape must be preserved after reset" + assert (drift_after >= lo - 1e-6).all() and (drift_after <= hi + 1e-6).all(), ( + f"Drift after reset must be in [{lo}, {hi}], got [{drift_after.min():.4f}, {drift_after.max():.4f}]" + ) + assert not torch.allclose(drift_after, drift_before), ( + "reset() must resample drift; values must change from initial sample" + ) diff --git a/source/isaaclab/test/sensors/test_update_ray_caster_kernel.py b/source/isaaclab/test/sensors/test_update_ray_caster_kernel.py new file mode 100644 index 000000000000..65402518b7a0 --- /dev/null +++ b/source/isaaclab/test/sensors/test_update_ray_caster_kernel.py @@ -0,0 +1,510 @@ +# 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 + +"""Unit tests for :func:`update_ray_caster_kernel`. + +These tests exercise the kernel directly with hand-crafted warp arrays and +analytically computed expected outputs. No simulation, no stage, no AppLauncher +— just warp on CPU (or CUDA when available). +""" + +from __future__ import annotations + +import importlib.util +import math +import os + +import numpy as np +import pytest +import torch +import warp as wp + +# Import the kernel module directly to avoid pulling in the full isaaclab package +# (which requires Isaac Sim / Omniverse dependencies). The kernel file itself only +# depends on warp. +_KERNEL_PATH = os.path.join( + os.path.dirname(__file__), + os.pardir, + os.pardir, + "isaaclab", + "sensors", + "ray_caster", + "kernels.py", +) +_spec = importlib.util.spec_from_file_location("ray_caster_kernels", os.path.normpath(_KERNEL_PATH)) +_mod = importlib.util.module_from_spec(_spec) +_spec.loader.exec_module(_mod) + +update_ray_caster_kernel = _mod.update_ray_caster_kernel +ALIGNMENT_WORLD = _mod.ALIGNMENT_WORLD +ALIGNMENT_YAW = _mod.ALIGNMENT_YAW +ALIGNMENT_BASE = _mod.ALIGNMENT_BASE + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +wp.init() +DEVICE = "cuda:0" if wp.is_cuda_available() else "cpu" +TORCH_DEVICE = torch.device(DEVICE) +ATOL = 1e-5 + + +def _make_transform(pos: tuple[float, float, float], quat_xyzw: tuple[float, float, float, float]) -> wp.array: + """Create a warp transformf array (1,) from position and xyzw quaternion.""" + t = torch.tensor([[pos[0], pos[1], pos[2], quat_xyzw[0], quat_xyzw[1], quat_xyzw[2], quat_xyzw[3]]], device=DEVICE) + return wp.from_torch(t.contiguous()).view(wp.transformf) + + +def _identity_quat() -> tuple[float, float, float, float]: + """Return identity quaternion in xyzw.""" + return (0.0, 0.0, 0.0, 1.0) + + +def _yaw_quat(yaw_rad: float) -> tuple[float, float, float, float]: + """Pure yaw quaternion in xyzw.""" + return (0.0, 0.0, math.sin(yaw_rad / 2), math.cos(yaw_rad / 2)) + + +def _euler_to_quat_xyzw(roll: float, pitch: float, yaw: float) -> tuple[float, float, float, float]: + """Euler angles (intrinsic XYZ) to quaternion in xyzw convention.""" + q = torch.zeros(1, 4) + cr, sr = math.cos(roll / 2), math.sin(roll / 2) + cp, sp = math.cos(pitch / 2), math.sin(pitch / 2) + cy, sy = math.cos(yaw / 2), math.sin(yaw / 2) + # xyzw + q[0, 0] = sr * cp * cy - cr * sp * sy + q[0, 1] = cr * sp * cy + sr * cp * sy + q[0, 2] = cr * cp * sy - sr * sp * cy + q[0, 3] = cr * cp * cy + sr * sp * sy + return tuple(q[0].tolist()) + + +def _quat_rotate(q_xyzw: tuple, v: tuple) -> np.ndarray: + """Rotate vector v by quaternion q (xyzw) using numpy.""" + qx, qy, qz, qw = q_xyzw + # quaternion rotation: v' = q * v * q^-1 + # Using the formula: v' = v + 2*w*(w×v) + 2*(q_vec × (q_vec × v + w*v)) + # Simpler: v' = v + 2w(q×v) + 2(q×(q×v)) + q_vec = np.array([qx, qy, qz]) + v = np.array(v) + t = 2.0 * np.cross(q_vec, v) + return v + qw * t + np.cross(q_vec, t) + + +def _launch_kernel( + transforms: wp.array, + env_mask: wp.array, + offset_pos: wp.array, + offset_quat: wp.array, + drift: wp.array, + ray_cast_drift: wp.array, + ray_starts_local: wp.array, + ray_directions_local: wp.array, + alignment_mode: int, + num_envs: int, + num_rays: int, +) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: + """Launch the kernel and return (pos_w, quat_w, ray_starts_w, ray_directions_w) as numpy arrays.""" + pos_w = wp.zeros(num_envs, dtype=wp.vec3f, device=DEVICE) + quat_w = wp.zeros(num_envs, dtype=wp.quatf, device=DEVICE) + ray_starts_w = wp.zeros((num_envs, num_rays), dtype=wp.vec3f, device=DEVICE) + ray_directions_w = wp.zeros((num_envs, num_rays), dtype=wp.vec3f, device=DEVICE) + + wp.launch( + update_ray_caster_kernel, + dim=(num_envs, num_rays), + inputs=[ + transforms, + env_mask, + offset_pos, + offset_quat, + drift, + ray_cast_drift, + ray_starts_local, + ray_directions_local, + alignment_mode, + ], + outputs=[pos_w, quat_w, ray_starts_w, ray_directions_w], + device=DEVICE, + ) + wp.synchronize_device(DEVICE) + + return ( + wp.to_torch(pos_w).cpu().numpy(), + wp.to_torch(quat_w).cpu().numpy(), + wp.to_torch(ray_starts_w).cpu().numpy(), + wp.to_torch(ray_directions_w).cpu().numpy(), + ) + + +def _make_inputs( + view_pos=(0.0, 0.0, 0.0), + view_quat=None, + offset_pos=(0.0, 0.0, 0.0), + offset_quat=None, + drift=(0.0, 0.0, 0.0), + ray_cast_drift=(0.0, 0.0, 0.0), + ray_start=(0.0, 0.0, 0.0), + ray_dir=(0.0, 0.0, -1.0), + num_envs=1, +): + """Build all kernel input arrays for a single-ray, single (or multi)-env scenario.""" + if view_quat is None: + view_quat = _identity_quat() + if offset_quat is None: + offset_quat = _identity_quat() + + transforms = _make_transform(view_pos, view_quat) + if num_envs > 1: + # Replicate the same transform for all envs + t_torch = wp.to_torch(transforms).repeat(num_envs, 1) + transforms = wp.from_torch(t_torch.contiguous()).view(wp.transformf) + + mask_t = torch.ones(num_envs, dtype=torch.bool, device=TORCH_DEVICE) + env_mask = wp.from_torch(mask_t) + + op = torch.tensor( + [[offset_pos[0], offset_pos[1], offset_pos[2]]] * num_envs, dtype=torch.float32, device=TORCH_DEVICE + ) + offset_pos_wp = wp.from_torch(op.contiguous(), dtype=wp.vec3f) + + oq = torch.tensor( + [[offset_quat[0], offset_quat[1], offset_quat[2], offset_quat[3]]] * num_envs, + dtype=torch.float32, + device=TORCH_DEVICE, + ) + offset_quat_wp = wp.from_torch(oq.contiguous(), dtype=wp.quatf) + + d = torch.tensor([[drift[0], drift[1], drift[2]]] * num_envs, dtype=torch.float32, device=TORCH_DEVICE) + drift_wp = wp.from_torch(d.contiguous(), dtype=wp.vec3f) + + rcd = torch.tensor( + [[ray_cast_drift[0], ray_cast_drift[1], ray_cast_drift[2]]] * num_envs, dtype=torch.float32, device=TORCH_DEVICE + ) + rcd_wp = wp.from_torch(rcd.contiguous(), dtype=wp.vec3f) + + rs = torch.tensor( + [[[ray_start[0], ray_start[1], ray_start[2]]]] * num_envs, dtype=torch.float32, device=TORCH_DEVICE + ) + rs_wp = wp.from_torch(rs.contiguous(), dtype=wp.vec3f) + + rd = torch.tensor([[[ray_dir[0], ray_dir[1], ray_dir[2]]]] * num_envs, dtype=torch.float32, device=TORCH_DEVICE) + rd_wp = wp.from_torch(rd.contiguous(), dtype=wp.vec3f) + + return transforms, env_mask, offset_pos_wp, offset_quat_wp, drift_wp, rcd_wp, rs_wp, rd_wp + + +# --------------------------------------------------------------------------- +# Tests +# --------------------------------------------------------------------------- + + +class TestUpdateRayCasterKernel: + """Unit tests for update_ray_caster_kernel launched directly with warp arrays.""" + + def test_identity_passthrough(self): + """All identity/zero inputs → pos_w = origin, quat_w = identity, rays unchanged.""" + inputs = _make_inputs(ray_start=(1.0, 2.0, 3.0), ray_dir=(0.0, 0.0, -1.0)) + pos_w, quat_w, starts_w, dirs_w = _launch_kernel(*inputs, alignment_mode=0, num_envs=1, num_rays=1) + + np.testing.assert_allclose(pos_w[0], [0, 0, 0], atol=ATOL) + np.testing.assert_allclose(quat_w[0], [0, 0, 0, 1], atol=ATOL) + # World mode, identity: ray_start_w = local_start + pos (= local_start + origin) + np.testing.assert_allclose(starts_w[0, 0], [1, 2, 3], atol=ATOL) + np.testing.assert_allclose(dirs_w[0, 0], [0, 0, -1], atol=ATOL) + + # Same for yaw and base — all should agree at identity + for mode in [1, 2]: + inputs = _make_inputs(ray_start=(1.0, 2.0, 3.0), ray_dir=(0.0, 0.0, -1.0)) + _, _, starts_w2, dirs_w2 = _launch_kernel(*inputs, alignment_mode=mode, num_envs=1, num_rays=1) + np.testing.assert_allclose(starts_w2[0, 0], [1, 2, 3], atol=ATOL) + np.testing.assert_allclose(dirs_w2[0, 0], [0, 0, -1], atol=ATOL) + + def test_offset_composition(self): + """View at (1,0,2) yawed 90° + offset (0,1,0) → combined_pos = (1,-1,2). + + 90° yaw: quat = (0, 0, sin(45°), cos(45°)) = (0, 0, 0.7071, 0.7071) + quat_rotate(90°yaw, (0,1,0)) = (-1, 0, 0) [Y axis maps to -X] + combined_pos = (1,0,2) + (-1,0,0) = (0,0,2) + combined_quat = yaw90 * identity = yaw90 + """ + yaw90 = _yaw_quat(math.pi / 2) + inputs = _make_inputs( + view_pos=(1.0, 0.0, 2.0), + view_quat=yaw90, + offset_pos=(0.0, 1.0, 0.0), + ) + pos_w, quat_w, _, _ = _launch_kernel(*inputs, alignment_mode=2, num_envs=1, num_rays=1) + + expected_offset_rotated = _quat_rotate(yaw90, (0, 1, 0)) # (-1, 0, 0) + expected_pos = np.array([1, 0, 2]) + expected_offset_rotated + np.testing.assert_allclose(pos_w[0], expected_pos, atol=ATOL) + np.testing.assert_allclose(quat_w[0], list(yaw90), atol=ATOL) + + def test_world_alignment_ignores_rotation(self): + """World mode: ray starts = local_start + combined_pos, directions unchanged. + + Sensor at (0,0,5), pitched 45° around Y. Local ray at (+1,0,0), direction (0,0,-1). + World mode should NOT rotate the ray start or direction. + """ + pitch45 = _euler_to_quat_xyzw(0, math.pi / 4, 0) + inputs = _make_inputs( + view_pos=(0.0, 0.0, 5.0), + view_quat=pitch45, + ray_start=(1.0, 0.0, 0.0), + ray_dir=(0.0, 0.0, -1.0), + ) + pos_w, _, starts_w, dirs_w = _launch_kernel(*inputs, alignment_mode=0, num_envs=1, num_rays=1) + + # ray_start_w = local_start + combined_pos = (1,0,0) + (0,0,5) = (1,0,5) + np.testing.assert_allclose(starts_w[0, 0], [1, 0, 5], atol=ATOL) + # direction unchanged + np.testing.assert_allclose(dirs_w[0, 0], [0, 0, -1], atol=ATOL) + + def test_yaw_alignment_rotates_starts_only(self): + """Yaw mode: ray starts rotated by yaw-only quaternion, directions unchanged. + + Sensor yawed 90° + pitched 30°. Local ray start at (+1, 0, 0). + Yaw-only extracts 90° yaw → rotates (+1,0,0) to (0,+1,0). + Direction (0,0,-1) is NOT rotated in yaw mode. + """ + q = _euler_to_quat_xyzw(0, math.pi / 6, math.pi / 2) # pitch 30°, yaw 90° + inputs = _make_inputs( + view_pos=(0.0, 0.0, 3.0), + view_quat=q, + ray_start=(1.0, 0.0, 0.0), + ray_dir=(0.0, 0.0, -1.0), + ) + pos_w, _, starts_w, dirs_w = _launch_kernel(*inputs, alignment_mode=1, num_envs=1, num_rays=1) + + # yaw-only of 90° yaw + 30° pitch → pure 90° yaw + yaw_only = _yaw_quat(math.pi / 2) + rotated_start = _quat_rotate(yaw_only, (1, 0, 0)) # (0, 1, 0) + expected_start = rotated_start + np.array([0, 0, 3]) # + combined_pos + np.testing.assert_allclose(starts_w[0, 0], expected_start, atol=ATOL) + + # direction unchanged in yaw mode + np.testing.assert_allclose(dirs_w[0, 0], [0, 0, -1], atol=ATOL) + + def test_base_alignment_rotates_starts_and_directions(self): + """Base mode: both ray starts and directions rotated by full combined quaternion. + + Sensor yawed 90°. Local ray at (+1, 0, 0), direction (0, 0, -1). + 90° yaw rotates: + (+1,0,0) → (0,+1,0) + (0,0,-1) → (0,0,-1) [yaw doesn't affect Z-down] + """ + yaw90 = _yaw_quat(math.pi / 2) + inputs = _make_inputs( + view_pos=(0.0, 0.0, 4.0), + view_quat=yaw90, + ray_start=(1.0, 0.0, 0.0), + ray_dir=(0.0, 0.0, -1.0), + ) + _, _, starts_w, dirs_w = _launch_kernel(*inputs, alignment_mode=2, num_envs=1, num_rays=1) + + rotated_start = _quat_rotate(yaw90, (1, 0, 0)) # (0, 1, 0) + expected_start = rotated_start + np.array([0, 0, 4]) + np.testing.assert_allclose(starts_w[0, 0], expected_start, atol=ATOL) + + rotated_dir = _quat_rotate(yaw90, (0, 0, -1)) # (0, 0, -1) — Z unaffected by yaw + np.testing.assert_allclose(dirs_w[0, 0], rotated_dir, atol=ATOL) + + def test_base_alignment_with_pitch_rotates_direction(self): + """Base mode with pitch: direction is rotated by the full orientation. + + Sensor pitched 90° around Y (looking forward instead of down). + Direction (0,0,-1) rotated by 90° pitch around Y → (-1,0,0). + """ + pitch90 = _euler_to_quat_xyzw(0, math.pi / 2, 0) + inputs = _make_inputs( + view_pos=(0.0, 0.0, 2.0), + view_quat=pitch90, + ray_start=(0.0, 0.0, 0.0), + ray_dir=(0.0, 0.0, -1.0), + ) + _, _, _, dirs_w = _launch_kernel(*inputs, alignment_mode=2, num_envs=1, num_rays=1) + + rotated_dir = _quat_rotate(pitch90, (0, 0, -1)) # (-1, 0, 0) + np.testing.assert_allclose(dirs_w[0, 0], rotated_dir, atol=ATOL) + + def test_ray_cast_drift_world_mode(self): + """World mode: ray_cast_drift XY is added raw to position, Z is NOT applied. + + drift = (0.5, 0.3, 0.7). In world mode: + pos_drifted = (combined_pos.x + 0.5, combined_pos.y + 0.3, combined_pos.z) + Note: Z component of ray_cast_drift is NOT added to position in any mode. + """ + inputs = _make_inputs( + view_pos=(1.0, 2.0, 3.0), + ray_cast_drift=(0.5, 0.3, 0.7), + ray_start=(0.0, 0.0, 0.0), + ray_dir=(0.0, 0.0, -1.0), + ) + _, _, starts_w, dirs_w = _launch_kernel(*inputs, alignment_mode=0, num_envs=1, num_rays=1) + + # World mode: pos_drifted = (1+0.5, 2+0.3, 3) = (1.5, 2.3, 3) + # ray_start_w = local_start + pos_drifted = (0,0,0) + (1.5, 2.3, 3) + np.testing.assert_allclose(starts_w[0, 0], [1.5, 2.3, 3.0], atol=ATOL) + np.testing.assert_allclose(dirs_w[0, 0], [0, 0, -1], atol=ATOL) + + def test_ray_cast_drift_yaw_mode(self): + """Yaw mode: ray_cast_drift XY is rotated by yaw-only quat, Z is NOT applied. + + Sensor yawed 90°, drift = (1.0, 0.0, 0.5). + yaw-rotated drift = quat_rotate(yaw90, (1,0,0.5)) — but only XY of the result + is used for pos_drifted. Actually looking at the kernel: + rot_drift = quat_rotate(yaw_q, rcd) # full rotation of the drift vector + pos_drifted = (combined_pos.x + rot_drift.x, combined_pos.y + rot_drift.y, combined_pos.z) + So the drift vector is fully rotated, but only XY of the result is added. + """ + yaw90 = _yaw_quat(math.pi / 2) + inputs = _make_inputs( + view_pos=(0.0, 0.0, 5.0), + view_quat=yaw90, + ray_cast_drift=(1.0, 0.0, 0.5), + ray_start=(0.0, 0.0, 0.0), + ray_dir=(0.0, 0.0, -1.0), + ) + _, _, starts_w, _ = _launch_kernel(*inputs, alignment_mode=1, num_envs=1, num_rays=1) + + # yaw90 rotates (1, 0, 0.5) → (0, 1, 0.5) [X→Y under 90° yaw, Z unchanged] + rot_drift = _quat_rotate(yaw90, (1, 0, 0.5)) + # pos_drifted = (0 + rot_drift.x, 0 + rot_drift.y, 5) — Z from combined_pos + expected_start = np.array([rot_drift[0], rot_drift[1], 5.0]) + # local_start = (0,0,0), rotated by yaw_q → still (0,0,0) + np.testing.assert_allclose(starts_w[0, 0], expected_start, atol=ATOL) + + def test_ray_cast_drift_base_mode(self): + """Base mode: ray_cast_drift XY is rotated by full combined_quat, Z is NOT applied. + + Sensor pitched 90° around Y, drift = (1.0, 0.0, 0.0). + Full rotation of (1,0,0) by 90° pitch around Y → (0, 0, -1). + pos_drifted = (combined_pos.x + 0, combined_pos.y + 0, combined_pos.z) — both XY of + rotated drift happen to be 0 in this case. + """ + pitch90 = _euler_to_quat_xyzw(0, math.pi / 2, 0) + inputs = _make_inputs( + view_pos=(0.0, 0.0, 5.0), + view_quat=pitch90, + ray_cast_drift=(1.0, 0.0, 0.0), + ray_start=(0.0, 0.0, 0.0), + ray_dir=(0.0, 0.0, -1.0), + ) + _, _, starts_w, _ = _launch_kernel(*inputs, alignment_mode=2, num_envs=1, num_rays=1) + + rot_drift = _quat_rotate(pitch90, (1, 0, 0)) # (0, 0, -1) + # pos_drifted = (0 + rot_drift.x, 0 + rot_drift.y, 5) = (0, 0, 5) + # local_start (0,0,0) rotated by pitch90 → still (0,0,0) + expected_start = np.array([rot_drift[0], rot_drift[1], 5.0]) + np.testing.assert_allclose(starts_w[0, 0], expected_start, atol=ATOL) + + def test_env_mask_skips_masked_envs(self): + """Masked-out environments retain sentinel values in output buffers. + + 2 envs, env 0 masked out (False), env 1 active (True). + Output buffers are pre-filled with sentinel (999). After kernel launch, + env 0 should still have 999, env 1 should have computed values. + """ + yaw90 = _yaw_quat(math.pi / 2) + + # Build transforms for 2 envs: both at (0,0,2) with yaw90 + t_single = torch.tensor( + [[0, 0, 2, yaw90[0], yaw90[1], yaw90[2], yaw90[3]]], + dtype=torch.float32, + device=DEVICE, + ) + t_both = t_single.repeat(2, 1).contiguous() + transforms = wp.from_torch(t_both).view(wp.transformf) + + # Mask: env 0 = False, env 1 = True + mask_t = torch.tensor([False, True], dtype=torch.bool, device=TORCH_DEVICE) + env_mask = wp.from_torch(mask_t) + + # Zero offsets and drifts for both envs + zero3 = torch.zeros(2, 3, dtype=torch.float32, device=TORCH_DEVICE) + offset_pos_wp = wp.from_torch(zero3.clone().contiguous(), dtype=wp.vec3f) + iq = torch.tensor([[0, 0, 0, 1]] * 2, dtype=torch.float32, device=TORCH_DEVICE) + offset_quat_wp = wp.from_torch(iq.contiguous(), dtype=wp.quatf) + drift_wp = wp.from_torch(zero3.clone().contiguous(), dtype=wp.vec3f) + rcd_wp = wp.from_torch(zero3.clone().contiguous(), dtype=wp.vec3f) + + # Single ray per env + rs = torch.tensor([[[1, 0, 0]]] * 2, dtype=torch.float32, device=TORCH_DEVICE) + rs_wp = wp.from_torch(rs.contiguous(), dtype=wp.vec3f) + rd = torch.tensor([[[0, 0, -1]]] * 2, dtype=torch.float32, device=TORCH_DEVICE) + rd_wp = wp.from_torch(rd.contiguous(), dtype=wp.vec3f) + + # Pre-fill outputs with sentinel + sentinel = 999.0 + pos_w_t = torch.full((2, 3), sentinel, dtype=torch.float32, device=TORCH_DEVICE) + pos_w = wp.from_torch(pos_w_t.contiguous(), dtype=wp.vec3f) + quat_w_t = torch.full((2, 4), sentinel, dtype=torch.float32, device=TORCH_DEVICE) + quat_w = wp.from_torch(quat_w_t.contiguous(), dtype=wp.quatf) + starts_w_t = torch.full((2, 1, 3), sentinel, dtype=torch.float32, device=TORCH_DEVICE) + starts_w = wp.from_torch(starts_w_t.contiguous(), dtype=wp.vec3f) + dirs_w_t = torch.full((2, 1, 3), sentinel, dtype=torch.float32, device=TORCH_DEVICE) + dirs_w = wp.from_torch(dirs_w_t.contiguous(), dtype=wp.vec3f) + + wp.launch( + update_ray_caster_kernel, + dim=(2, 1), + inputs=[transforms, env_mask, offset_pos_wp, offset_quat_wp, drift_wp, rcd_wp, rs_wp, rd_wp, 2], + outputs=[pos_w, quat_w, starts_w, dirs_w], + device=DEVICE, + ) + wp.synchronize_device(DEVICE) + + pos_np = wp.to_torch(pos_w).cpu().numpy() + quat_np = wp.to_torch(quat_w).cpu().numpy() + starts_np = wp.to_torch(starts_w).cpu().numpy() + dirs_np = wp.to_torch(dirs_w).cpu().numpy() + + # Env 0 (masked): all outputs should still be sentinel + np.testing.assert_allclose(pos_np[0], [sentinel] * 3, atol=ATOL) + np.testing.assert_allclose(quat_np[0], [sentinel] * 4, atol=ATOL) + np.testing.assert_allclose(starts_np[0, 0], [sentinel] * 3, atol=ATOL) + np.testing.assert_allclose(dirs_np[0, 0], [sentinel] * 3, atol=ATOL) + + # Env 1 (active): should have computed values + np.testing.assert_allclose(pos_np[1], [0, 0, 2], atol=ATOL) + np.testing.assert_allclose(quat_np[1], list(yaw90), atol=ATOL) + # Base mode: (1,0,0) rotated by yaw90 = (0,1,0), + pos (0,0,2) + expected_start = _quat_rotate(yaw90, (1, 0, 0)) + np.array([0, 0, 2]) + np.testing.assert_allclose(starts_np[1, 0], expected_start, atol=ATOL) + expected_dir = _quat_rotate(yaw90, (0, 0, -1)) # (0, 0, -1) unaffected by yaw + np.testing.assert_allclose(dirs_np[1, 0], expected_dir, atol=ATOL) + + def test_positional_drift_added_before_alignment(self): + """The `drift` parameter is added to combined_pos before ray transformation. + + Verify that drift shifts the sensor position (and therefore ray starts) + equally across all alignment modes. + """ + drift_val = (0.0, 0.0, 1.5) # shift up 1.5m + results = {} + for mode_name, mode_int in [("world", 0), ("yaw", 1), ("base", 2)]: + inputs = _make_inputs( + view_pos=(0.0, 0.0, 3.0), + drift=drift_val, + ray_start=(0.0, 0.0, 0.0), + ray_dir=(0.0, 0.0, -1.0), + ) + pos_w, _, starts_w, _ = _launch_kernel(*inputs, alignment_mode=mode_int, num_envs=1, num_rays=1) + results[mode_name] = (pos_w, starts_w) + + # All modes: pos_w should be (0, 0, 4.5) = view_pos + drift + for mode_name in ["world", "yaw", "base"]: + np.testing.assert_allclose( + results[mode_name][0][0], + [0, 0, 4.5], + atol=ATOL, + err_msg=f"{mode_name} mode: pos_w should include drift", + ) + # ray_start_w Z should also reflect the drifted position + assert results[mode_name][1][0, 0, 2] == pytest.approx(4.5, abs=ATOL), ( + f"{mode_name} mode: ray start Z should be 4.5" + ) 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_physx_scene_data_provider_visualizer_contract.py b/source/isaaclab/test/sim/test_physx_scene_data_provider_visualizer_contract.py index abbc3046655a..927fe351d202 100644 --- a/source/isaaclab/test/sim/test_physx_scene_data_provider_visualizer_contract.py +++ b/source/isaaclab/test/sim/test_physx_scene_data_provider_visualizer_contract.py @@ -8,6 +8,7 @@ from __future__ import annotations from types import SimpleNamespace +from unittest.mock import patch from isaaclab_physx.scene_data_providers import PhysxSceneDataProvider @@ -15,47 +16,21 @@ def _make_provider(): - provider = object.__new__(PhysxSceneDataProvider) - provider._force_usd_fallback_for_newton_model_build = False - return provider + return object.__new__(PhysxSceneDataProvider) -def test_get_newton_model_for_env_ids_builds_and_caches_sorted_keys(): +def test_get_newton_model_returns_model_when_sync_enabled(): + """Callers receive the full Newton model from :meth:`get_newton_model`.""" provider = _make_provider() provider._needs_newton_sync = True provider._newton_model = "full-model" - provider._filtered_newton_model = None - provider._filtered_env_ids_key = None - build_calls = [] + assert provider.get_newton_model() == "full-model" - def _fake_build(env_ids): - build_calls.append(env_ids) - provider._filtered_newton_model = f"filtered-{env_ids}" - provider._build_filtered_newton_model = _fake_build - - # None asks for the full model. - assert provider.get_newton_model_for_env_ids(None) == "full-model" - - # First subset request builds using sorted env id key. - model_a = provider.get_newton_model_for_env_ids([3, 1]) - assert model_a == "filtered-[1, 3]" - assert build_calls == [[1, 3]] - - # Equivalent request should use cache and not rebuild. - model_b = provider.get_newton_model_for_env_ids([1, 3]) - assert model_b == "filtered-[1, 3]" - assert build_calls == [[1, 3]] - - # Different subset rebuilds. - model_c = provider.get_newton_model_for_env_ids([2]) - assert model_c == "filtered-[2]" - assert build_calls == [[1, 3], [2]] - - -def test_try_use_prebuilt_artifact_populates_provider_state(): - """Provider should consume scene-time prebuilt artifact as fast path.""" +@patch("isaaclab_physx.scene_data_providers.physx_scene_data_provider.replace_newton_shape_colors", lambda m, s: None) +def test_load_prebuilt_artifact_populates_provider_state(): + """Loading the prebuilt artifact sets model, state, and rigid-body paths.""" provider = _make_provider() artifact = VisualizerPrebuiltArtifacts( model="prebuilt-model", @@ -65,6 +40,7 @@ def test_try_use_prebuilt_artifact_populates_provider_state(): num_envs=4, ) provider._simulation_context = SimpleNamespace(get_scene_data_visualizer_prebuilt_artifact=lambda: artifact) + provider._stage = None provider._xform_views = {"old": object()} provider._view_body_index_map = {"old": [1]} @@ -74,19 +50,13 @@ def test_try_use_prebuilt_artifact_populates_provider_state(): provider._orientations_buf = object() provider._covered_buf = object() provider._xform_mask_buf = object() - provider._env_id_to_body_indices = {0: [0]} - provider._filtered_newton_model = "old-filtered-model" - provider._filtered_newton_state = "old-filtered-state" - provider._filtered_env_ids_key = (0,) - provider._filtered_body_indices = [0] - provider._stage = None - - assert provider._try_use_prebuilt_newton_artifact() is True + provider._load_newton_model_from_prebuilt_artifact() assert provider._newton_model == "prebuilt-model" assert provider._newton_state == "prebuilt-state" assert provider._rigid_body_paths == ["/World/envs/env_0/A"] assert provider._rigid_body_view_paths == ["/World/envs/env_0/A", "/World/envs/env_0/Robot"] assert provider._num_envs_at_last_newton_build == 4 + assert provider._last_newton_model_build_source == "prebuilt" assert provider._xform_views == {} assert provider._view_body_index_map == {} assert provider._view_order_tensors == {} @@ -95,32 +65,13 @@ def test_try_use_prebuilt_artifact_populates_provider_state(): assert provider._orientations_buf is None assert provider._covered_buf is None assert provider._xform_mask_buf is None - assert provider._env_id_to_body_indices == {} - assert provider._filtered_newton_model is None - assert provider._filtered_newton_state is None - assert provider._filtered_env_ids_key is None - assert provider._filtered_body_indices == [] -def test_try_use_prebuilt_artifact_respects_force_usd_fallback_flag(): - """Force flag should disable prebuilt fast path even when artifact is available.""" +def test_load_prebuilt_artifact_missing_sets_error_state(): + """When no artifact is registered, model/state stay unset.""" provider = _make_provider() - provider._force_usd_fallback_for_newton_model_build = True - artifact = VisualizerPrebuiltArtifacts( - model="prebuilt-model", - state="prebuilt-state", - rigid_body_paths=["/World/envs/env_0/A"], - articulation_paths=["/World/envs/env_0/Robot"], - num_envs=4, - ) - provider._simulation_context = SimpleNamespace(get_scene_data_visualizer_prebuilt_artifact=lambda: artifact) - - assert provider._try_use_prebuilt_newton_artifact() is False - - -def test_build_newton_model_from_usd_short_circuits_when_prebuilt_available(): - """If prebuilt artifact is available, USD fallback should not run.""" - provider = _make_provider() - provider._try_use_prebuilt_newton_artifact = lambda: True - provider._build_newton_model_from_usd() - assert provider._last_newton_model_build_source == "prebuilt" + provider._simulation_context = SimpleNamespace(get_scene_data_visualizer_prebuilt_artifact=lambda: None) + provider._load_newton_model_from_prebuilt_artifact() + assert provider._last_newton_model_build_source == "missing" + assert provider._newton_model is None + assert provider._newton_state is None diff --git a/source/isaaclab/test/sim/test_simulation_context.py b/source/isaaclab/test/sim/test_simulation_context.py index c03413838e3e..6ea578a85e30 100644 --- a/source/isaaclab/test/sim/test_simulation_context.py +++ b/source/isaaclab/test/sim/test_simulation_context.py @@ -290,6 +290,86 @@ def test_render(): assert sim.is_playing() +@pytest.mark.isaacsim_ci +def test_render_pumps_app_update_without_visualizer(): + """Regression test for issue #5052: headless video must pump Kit when no visualizer does. + + Originally ``SimulationContext.render()`` called ``omni.kit.app.get_app().update()`` when + no visualizer had ``pumps_app_update()`` (see PR #5056). The same contract is now implemented + from :func:`~isaaclab.envs.utils.recording_hooks.run_recording_hooks_after_visualizers`, which calls + :func:`~isaaclab_physx.renderers.isaac_rtx_renderer_utils.pump_kit_app_for_headless_video_render_if_needed` + when ``/isaaclab/video/enabled`` is set (as with ``--video``), which in turn calls + ``ensure_isaac_rtx_render_update()`` (guarded by ``is_rendering`` and the no-pumping-visualizer check). + + Without this path, replicator render products used for ``rgb_array`` / RecordVideo stay stale (black frames). + """ + from unittest.mock import MagicMock, patch + + cfg = SimulationCfg(dt=0.01) + sim = SimulationContext(cfg) + sim.reset() + + sim.set_setting("/isaaclab/video/enabled", True) + sim.set_setting("/isaaclab/render/rtx_sensors", True) + + mock_app = MagicMock() + mock_app.is_running.return_value = True + + with ( + patch("isaaclab.utils.version.has_kit", return_value=True), + patch( + "isaaclab_physx.renderers.isaac_rtx_renderer_utils._get_stage_streaming_busy", + return_value=False, + ), + patch("omni.kit.app.get_app", return_value=mock_app), + ): + sim.render() + + mock_app.update.assert_called_once() + + +@pytest.mark.isaacsim_ci +def test_render_skips_app_update_when_visualizer_pumps_it(): + """Regression test: do not pump Kit in the headless-video path when a visualizer already does. + + A visualizer with ``pumps_app_update() == True`` (e.g. KitVisualizer) calls ``app.update()`` in + its own ``step()``. The recording-hook pump must then skip + ``ensure_isaac_rtx_render_update`` so we do not double-pump the Kit loop. + """ + from unittest.mock import MagicMock, patch + + from isaaclab.visualizers.base_visualizer import BaseVisualizer + + cfg = SimulationCfg(dt=0.01) + sim = SimulationContext(cfg) + sim.reset() + + sim.set_setting("/isaaclab/video/enabled", True) + sim.set_setting("/isaaclab/render/rtx_sensors", True) + + mock_viz = MagicMock(spec=BaseVisualizer) + mock_viz.pumps_app_update.return_value = True + mock_viz.is_closed = False + mock_viz.is_running.return_value = True + mock_viz.is_rendering_paused.return_value = False + mock_viz.is_training_paused.return_value = False + mock_viz.get_rendering_dt.return_value = None + sim._visualizers = [mock_viz] + + mock_app = MagicMock() + mock_app.is_running.return_value = True + + with ( + patch("isaaclab.utils.version.has_kit", return_value=True), + patch("omni.kit.app.get_app", return_value=mock_app), + ): + sim.render() + + mock_app.update.assert_not_called() + + sim._visualizers = [] + + """ Stage Operations Tests """ diff --git a/source/isaaclab/test/sim/test_simulation_context_visualizers.py b/source/isaaclab/test/sim/test_simulation_context_visualizers.py index f0a1294d2b4a..d5fa5ffbb2ce 100644 --- a/source/isaaclab/test/sim/test_simulation_context_visualizers.py +++ b/source/isaaclab/test/sim/test_simulation_context_visualizers.py @@ -30,11 +30,13 @@ class _FakeProvider: def __init__(self): self.update_calls = [] - def update(self, env_ids=None): - self.update_calls.append(env_ids) + def update(self): + self.update_calls.append(True) class _FakeVisualizer: + """Minimal visualizer for orchestration tests.""" + def __init__( self, *, @@ -45,6 +47,7 @@ def __init__( training_paused_steps=0, raises_on_step=False, requires_forward=False, + pumps_app_update=False, ): self._env_ids = env_ids self._running = running @@ -53,6 +56,7 @@ def __init__( self._training_paused_steps = training_paused_steps self._raises_on_step = raises_on_step self._requires_forward = requires_forward + self._pumps_app_update = pumps_app_update self.step_calls = [] self.close_calls = 0 @@ -87,6 +91,9 @@ def get_visualized_env_ids(self): def requires_forward_before_step(self): return self._requires_forward + def pumps_app_update(self): + return self._pumps_app_update + def _make_context(visualizers, provider=None): ctx = object.__new__(SimulationContext) @@ -97,7 +104,7 @@ def _make_context(visualizers, provider=None): return ctx -def test_update_scene_data_provider_unions_env_ids_and_forwards(): +def test_update_scene_data_provider_forwards_and_updates_provider(): provider = _FakeProvider() viz_a = _FakeVisualizer(env_ids=[0, 2], requires_forward=True) viz_b = _FakeVisualizer(env_ids=[2, 3]) @@ -107,7 +114,7 @@ def test_update_scene_data_provider_unions_env_ids_and_forwards(): ctx.update_scene_data_provider() assert ctx.physics_manager.forward_calls == 1 - assert provider.update_calls == [[0, 2, 3]] + assert provider.update_calls == [True] assert ctx._visualizer_step_counter == 1 @@ -116,7 +123,7 @@ def test_update_scene_data_provider_force_forward_with_no_visualizers(): ctx = _make_context([], provider=provider) ctx.update_scene_data_provider(force_require_forward=True) assert ctx.physics_manager.forward_calls == 1 - assert provider.update_calls == [None] + assert provider.update_calls == [True] def test_update_visualizers_removes_closed_nonrunning_and_failed(caplog): @@ -136,10 +143,21 @@ def test_update_visualizers_removes_closed_nonrunning_and_failed(caplog): assert stopped_viz.close_calls == 1 assert failing_viz.close_calls == 1 assert paused_viz.close_calls == 0 + assert paused_viz.step_calls == [0.0] assert healthy_viz.step_calls == [0.1] assert any("Error stepping visualizer" in r.message for r in caplog.records) +def test_update_visualizers_skips_zero_dt_for_paused_app_pumping_visualizer(): + provider = _FakeProvider() + paused_app_pumping_viz = _FakeVisualizer(rendering_paused=True, pumps_app_update=True) + ctx = _make_context([paused_app_pumping_viz], provider=provider) + + ctx.update_visualizers(0.3) + + assert paused_app_pumping_viz.step_calls == [] + + def test_update_visualizers_handles_training_pause_loop(): provider = _FakeProvider() viz = _FakeVisualizer(training_paused_steps=1) @@ -161,9 +179,9 @@ def get_metadata(self) -> dict: def get_newton_model(self): return "dummy-model" - def get_newton_state(self, env_ids: list[int] | None): - self.state_calls.append(env_ids) - return {"state_call": len(self.state_calls), "env_ids": env_ids} + def get_newton_state(self): + self.state_calls.append(None) + return {"state_call": len(self.state_calls)} class _DummyViserViewer: @@ -203,20 +221,23 @@ def _fake_create_viewer(self, record_to_viser: str | None, metadata: dict | None assert visualizer._sim_time == pytest.approx(0.25) assert viewer.calls[0][0] == "begin_frame" assert viewer.calls[0][1] == pytest.approx(0.25) - assert viewer.calls[1] == ("log_state", {"state_call": 2, "env_ids": None}) + # log_state passes through get_newton_state() as-is; no env_ids (or other) keys are merged in. + assert viewer.calls[1] == ("log_state", {"state_call": 2}) assert viewer.calls[2] == ("end_frame",) @pytest.mark.parametrize( - ("cfg_max_worlds", "expected_max_worlds"), + ("cfg_max_visible_envs", "expected_visible"), [ (None, None), - (0, 0), - (3, 3), + (0, []), + (3, [0, 1, 2]), ], ) -def test_viser_visualizer_create_viewer_forwards_max_worlds( - monkeypatch: pytest.MonkeyPatch, cfg_max_worlds: int | None, expected_max_worlds: int | None +def test_viser_visualizer_create_viewer_applies_visible_worlds( + monkeypatch: pytest.MonkeyPatch, + cfg_max_visible_envs: int | None, + expected_visible: list[int] | None, ): captured = {} @@ -240,8 +261,11 @@ def __init__( "metadata": metadata, } - def set_model(self, model: Any, max_worlds: int | None) -> None: - captured["set_model"] = {"model": model, "max_worlds": max_worlds} + def set_model(self, model: Any) -> None: + captured["set_model"] = model + + def set_visible_worlds(self, worlds) -> None: + captured["visible_worlds"] = worlds def set_world_offsets(self, spacing) -> None: captured["set_world_offsets"] = tuple(spacing) @@ -254,25 +278,33 @@ def set_world_offsets(self, spacing) -> None: ) monkeypatch.setattr(viser_visualizer.ViserVisualizer, "_set_viser_camera_view", lambda self, pose: None) - cfg = ViserVisualizerCfg(max_worlds=cfg_max_worlds, open_browser=False) + cfg = ViserVisualizerCfg( + max_visible_envs=cfg_max_visible_envs, + open_browser=False, + randomly_sample_visible_envs=False, + ) visualizer = viser_visualizer.ViserVisualizer(cfg) visualizer._model = "dummy-model" + visualizer._env_ids = None # normally set by initialize() -> _compute_visualized_env_ids() visualizer._create_viewer(record_to_viser="record.viser", metadata={"num_envs": 8}) - assert captured["set_model"] == {"model": "dummy-model", "max_worlds": expected_max_worlds} + assert captured["set_model"] == "dummy-model" + assert captured["visible_worlds"] == expected_visible assert captured["set_world_offsets"] == (0.0, 0.0, 0.0) @pytest.mark.parametrize( - ("cfg_max_worlds", "expected_max_worlds"), + ("cfg_max_visible_envs", "expected_visible"), [ (None, None), - (0, 0), - (3, 3), + (0, []), + (3, [0, 1, 2]), ], ) -def test_rerun_visualizer_initialize_forwards_max_worlds_and_world_offsets( - monkeypatch: pytest.MonkeyPatch, cfg_max_worlds: int | None, expected_max_worlds: int | None +def test_rerun_visualizer_initialize_applies_visible_worlds_and_world_offsets( + monkeypatch: pytest.MonkeyPatch, + cfg_max_visible_envs: int | None, + expected_visible: list[int] | None, ): captured = {} @@ -300,8 +332,11 @@ def __init__( "record_to_rrd": record_to_rrd, } - def set_model(self, model: Any, max_worlds: int | None = None) -> None: - captured["set_model"] = {"model": model, "max_worlds": max_worlds} + def set_model(self, model: Any) -> None: + captured["set_model"] = model + + def set_visible_worlds(self, worlds) -> None: + captured["visible_worlds"] = worlds def set_world_offsets(self, spacing) -> None: captured["set_world_offsets"] = tuple(spacing) @@ -316,8 +351,8 @@ def get_metadata(self) -> dict: def get_newton_model(self): return "dummy-model" - def get_newton_state(self, env_ids: list[int] | None): - return {"env_ids": env_ids} + def get_newton_state(self): + return {"ok": True} monkeypatch.setattr(rerun_visualizer, "NewtonViewerRerun", _FakeNewtonViewerRerun) monkeypatch.setattr( @@ -331,11 +366,16 @@ def get_newton_state(self, env_ids: list[int] | None): ) monkeypatch.setattr(rerun_visualizer.RerunVisualizer, "_apply_camera_pose", lambda self, pose: None) - cfg = RerunVisualizerCfg(open_browser=False, max_worlds=cfg_max_worlds) + cfg = RerunVisualizerCfg( + open_browser=False, + max_visible_envs=cfg_max_visible_envs, + randomly_sample_visible_envs=False, + ) visualizer = rerun_visualizer.RerunVisualizer(cfg) visualizer.initialize(cast(Any, _DummyRerunSceneDataProvider())) - assert captured["set_model"] == {"model": "dummy-model", "max_worlds": expected_max_worlds} + assert captured["set_model"] == "dummy-model" + assert captured["visible_worlds"] == expected_visible assert captured["set_world_offsets"] == (0.0, 0.0, 0.0) @@ -398,6 +438,8 @@ def _make_context_with_settings( ctx._has_gui = has_gui ctx._has_offscreen_render = has_offscreen_render ctx._xr_enabled = False + ctx._pending_camera_view = None + ctx._render_generation = 0 ctx._visualizers = [] ctx._scene_data_provider = _FakeProvider() ctx._scene_data_requirements = None @@ -438,7 +480,7 @@ def test_explicit_unknown_visualizer_type_raises(): "/isaaclab/visualizer/types": "bogus_viz", "/isaaclab/visualizer/explicit": True, "/isaaclab/visualizer/disable_all": False, - "/isaaclab/visualizer/max_worlds": None, + "/isaaclab/visualizer/max_visible_envs": None, } ctx = _make_context_with_settings(settings) @@ -452,7 +494,7 @@ def test_explicit_missing_package_raises(monkeypatch: pytest.MonkeyPatch): "/isaaclab/visualizer/types": "rerun", "/isaaclab/visualizer/explicit": True, "/isaaclab/visualizer/disable_all": False, - "/isaaclab/visualizer/max_worlds": None, + "/isaaclab/visualizer/max_visible_envs": None, } ctx = _make_context_with_settings(settings) @@ -479,7 +521,7 @@ def test_explicit_visualizer_create_failure_raises(monkeypatch: pytest.MonkeyPat "/isaaclab/visualizer/types": "newton", "/isaaclab/visualizer/explicit": True, "/isaaclab/visualizer/disable_all": False, - "/isaaclab/visualizer/max_worlds": None, + "/isaaclab/visualizer/max_visible_envs": None, } ctx = _make_context_with_settings(settings, visualizer_cfgs=[failing_cfg]) @@ -498,7 +540,7 @@ def test_explicit_visualizer_init_failure_raises(monkeypatch: pytest.MonkeyPatch "/isaaclab/visualizer/types": "newton", "/isaaclab/visualizer/explicit": True, "/isaaclab/visualizer/disable_all": False, - "/isaaclab/visualizer/max_worlds": None, + "/isaaclab/visualizer/max_visible_envs": None, } ctx = _make_context_with_settings(settings, visualizer_cfgs=[failing_cfg]) @@ -516,7 +558,7 @@ def test_explicit_partial_valid_types_raises_for_invalid(): "/isaaclab/visualizer/types": "newton,bogus_viz", "/isaaclab/visualizer/explicit": True, "/isaaclab/visualizer/disable_all": False, - "/isaaclab/visualizer/max_worlds": None, + "/isaaclab/visualizer/max_visible_envs": None, } ctx = _make_context_with_settings(settings) @@ -530,7 +572,7 @@ def test_non_explicit_unknown_type_silently_skipped(caplog): "/isaaclab/visualizer/types": "bogus_viz", "/isaaclab/visualizer/explicit": False, "/isaaclab/visualizer/disable_all": False, - "/isaaclab/visualizer/max_worlds": None, + "/isaaclab/visualizer/max_visible_envs": None, } ctx = _make_context_with_settings(settings) @@ -546,7 +588,7 @@ def test_non_explicit_create_failure_silently_logged(monkeypatch: pytest.MonkeyP "/isaaclab/visualizer/types": "", "/isaaclab/visualizer/explicit": False, "/isaaclab/visualizer/disable_all": False, - "/isaaclab/visualizer/max_worlds": None, + "/isaaclab/visualizer/max_visible_envs": None, } ctx = _make_context_with_settings(settings, visualizer_cfgs=[failing_cfg]) 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/test/utils/test_string.py b/source/isaaclab/test/utils/test_string.py index ce443dec705b..22f51ab6f483 100644 --- a/source/isaaclab/test/utils/test_string.py +++ b/source/isaaclab/test/utils/test_string.py @@ -19,6 +19,7 @@ import pytest import isaaclab.utils.string as string_utils +from isaaclab.utils.string import _resolve_matching_names_impl def test_resolvable_string_metadata_is_non_eager(): @@ -251,3 +252,22 @@ def test_resolve_matching_names_values_with_basic_strings_and_preserved_order(): query_names = {"a|c": 1, "b": 0, "f": 2} with pytest.raises(ValueError): _ = string_utils.resolve_matching_names_values(query_names, target_names, preserve_order=True) + + +def test_clear_resolve_matching_names_cache(): + """Clearing the cache discards previously cached entries.""" + target_names = ["a", "b", "c"] + # Populate the cache + string_utils.resolve_matching_names("a", target_names) + info_before = _resolve_matching_names_impl.cache_info() + assert info_before.currsize > 0 + + # Clear the cache + string_utils.clear_resolve_matching_names_cache() + info_after = _resolve_matching_names_impl.cache_info() + assert info_after.currsize == 0 + + # Results are still correct after clearing + idx, names = string_utils.resolve_matching_names("a", target_names) + assert idx == [0] + assert names == ["a"] diff --git a/source/isaaclab/test/utils/test_wrench_composer.py b/source/isaaclab/test/utils/test_wrench_composer.py index 9b7c4aecaf15..37f6d5959aac 100644 --- a/source/isaaclab/test/utils/test_wrench_composer.py +++ b/source/isaaclab/test/utils/test_wrench_composer.py @@ -131,11 +131,13 @@ def test_wrench_composer_add_force(device: str, num_envs: int, num_bodies: int): ) forces = wp.from_numpy(forces_np, dtype=wp.vec3f, device=device) # Add forces to wrench composer - wrench_composer.add_forces_and_torques(forces=forces, body_ids=body_ids, env_ids=env_ids) + wrench_composer.add_forces_and_torques_index(forces=forces, body_ids=body_ids, env_ids=env_ids) # Add forces to hand-calculated composed force hand_calculated_composed_force_np[env_ids_np[:, None], body_ids_np[None, :], :] += forces_np + # Compose to body frame before checking output + wrench_composer.compose_to_body_frame() # Get composed force from wrench composer - composed_force_np = wrench_composer.composed_force.numpy() + composed_force_np = wrench_composer.out_force_b.numpy() assert np.allclose(composed_force_np, hand_calculated_composed_force_np, atol=1, rtol=1e-7) @@ -168,18 +170,20 @@ def test_wrench_composer_add_torque(device: str, num_envs: int, num_bodies: int) ) torques = wp.from_numpy(torques_np, dtype=wp.vec3f, device=device) # Add torques to wrench composer - wrench_composer.add_forces_and_torques(torques=torques, body_ids=body_ids, env_ids=env_ids) + wrench_composer.add_forces_and_torques_index(torques=torques, body_ids=body_ids, env_ids=env_ids) # Add torques to hand-calculated composed torque hand_calculated_composed_torque_np[env_ids_np[:, None], body_ids_np[None, :], :] += torques_np + # Compose to body frame before checking output + wrench_composer.compose_to_body_frame() # Get composed torque from wrench composer - composed_torque_np = wrench_composer.composed_torque.numpy() + composed_torque_np = wrench_composer.out_torque_b.numpy() assert np.allclose(composed_torque_np, hand_calculated_composed_torque_np, atol=1, rtol=1e-7) @pytest.mark.parametrize("device", ["cuda:0", "cpu"]) @pytest.mark.parametrize("num_envs", [1, 10, 100, 1000]) @pytest.mark.parametrize("num_bodies", [1, 3, 5, 10]) -def test_add_forces_at_positons(device: str, num_envs: int, num_bodies: int): +def test_add_forces_at_positions(device: str, num_envs: int, num_bodies: int): """Test adding forces at local positions (offset from link frame).""" rng = np.random.default_rng(seed=2) @@ -214,7 +218,7 @@ def test_add_forces_at_positons(device: str, num_envs: int, num_bodies: int): forces = wp.from_numpy(forces_np, dtype=wp.vec3f, device=device) positions = wp.from_numpy(positions_np, dtype=wp.vec3f, device=device) # Add forces at positions to wrench composer - wrench_composer.add_forces_and_torques( + wrench_composer.add_forces_and_torques_index( forces=forces, positions=positions, body_ids=body_ids, env_ids=env_ids ) # Add forces to hand-calculated composed force @@ -225,11 +229,13 @@ def test_add_forces_at_positons(device: str, num_envs: int, num_bodies: int): for j in range(num_bodies_np): hand_calculated_composed_torque_np[env_ids_np[i], body_ids_np[j], :] += torques_from_forces[i, j, :] + # Compose to body frame before checking output + wrench_composer.compose_to_body_frame() # Get composed force from wrench composer - composed_force_np = wrench_composer.composed_force.numpy() + composed_force_np = wrench_composer.out_force_b.numpy() assert np.allclose(composed_force_np, hand_calculated_composed_force_np, atol=1, rtol=1e-7) # Get composed torque from wrench composer - composed_torque_np = wrench_composer.composed_torque.numpy() + composed_torque_np = wrench_composer.out_torque_b.numpy() assert np.allclose(composed_torque_np, hand_calculated_composed_torque_np, atol=1, rtol=1e-7) @@ -267,13 +273,15 @@ def test_add_torques_at_position(device: str, num_envs: int, num_bodies: int): torques = wp.from_numpy(torques_np, dtype=wp.vec3f, device=device) positions = wp.from_numpy(positions_np, dtype=wp.vec3f, device=device) # Add torques at positions to wrench composer - wrench_composer.add_forces_and_torques( + wrench_composer.add_forces_and_torques_index( torques=torques, positions=positions, body_ids=body_ids, env_ids=env_ids ) # Add torques to hand-calculated composed torque hand_calculated_composed_torque_np[env_ids_np[:, None], body_ids_np[None, :], :] += torques_np + # Compose to body frame before checking output + wrench_composer.compose_to_body_frame() # Get composed torque from wrench composer - composed_torque_np = wrench_composer.composed_torque.numpy() + composed_torque_np = wrench_composer.out_torque_b.numpy() assert np.allclose(composed_torque_np, hand_calculated_composed_torque_np, atol=1, rtol=1e-7) @@ -319,7 +327,7 @@ def test_add_forces_and_torques_at_position(device: str, num_envs: int, num_bodi torques = wp.from_numpy(torques_np, dtype=wp.vec3f, device=device) positions = wp.from_numpy(positions_np, dtype=wp.vec3f, device=device) # Add forces and torques at positions to wrench composer - wrench_composer.add_forces_and_torques( + wrench_composer.add_forces_and_torques_index( forces=forces, torques=torques, positions=positions, body_ids=body_ids, env_ids=env_ids ) # Add forces to hand-calculated composed force @@ -330,11 +338,13 @@ def test_add_forces_and_torques_at_position(device: str, num_envs: int, num_bodi for j in range(num_bodies_np): hand_calculated_composed_torque_np[env_ids_np[i], body_ids_np[j], :] += torques_from_forces[i, j, :] hand_calculated_composed_torque_np[env_ids_np[:, None], body_ids_np[None, :], :] += torques_np + # Compose to body frame before checking output + wrench_composer.compose_to_body_frame() # Get composed force from wrench composer - composed_force_np = wrench_composer.composed_force.numpy() + composed_force_np = wrench_composer.out_force_b.numpy() assert np.allclose(composed_force_np, hand_calculated_composed_force_np, atol=1, rtol=1e-7) # Get composed torque from wrench composer - composed_torque_np = wrench_composer.composed_torque.numpy() + composed_torque_np = wrench_composer.out_torque_b.numpy() assert np.allclose(composed_torque_np, hand_calculated_composed_torque_np, atol=1, rtol=1e-7) @@ -368,14 +378,18 @@ def test_wrench_composer_reset(device: str, num_envs: int, num_bodies: int): forces = wp.from_numpy(forces_np, dtype=wp.vec3f, device=device) torques = wp.from_numpy(torques_np, dtype=wp.vec3f, device=device) # Add forces and torques to wrench composer - wrench_composer.add_forces_and_torques(forces=forces, torques=torques, body_ids=body_ids, env_ids=env_ids) + wrench_composer.add_forces_and_torques_index(forces=forces, torques=torques, body_ids=body_ids, env_ids=env_ids) # Reset wrench composer wrench_composer.reset() - # Get composed force and torque from wrench composer - composed_force_np = wrench_composer.composed_force.numpy() - composed_torque_np = wrench_composer.composed_torque.numpy() - assert np.allclose(composed_force_np, np.zeros((num_envs, num_bodies, 3)), atol=1, rtol=1e-7) - assert np.allclose(composed_torque_np, np.zeros((num_envs, num_bodies, 3)), atol=1, rtol=1e-7) + # Check all 7 buffers are zero (5 input + 2 output) + zeros = np.zeros((num_envs, num_bodies, 3), dtype=np.float32) + assert np.allclose(wrench_composer.global_force_w.numpy(), zeros, atol=1, rtol=1e-7) + assert np.allclose(wrench_composer.global_torque_w.numpy(), zeros, atol=1, rtol=1e-7) + assert np.allclose(wrench_composer.global_force_at_com_w.numpy(), zeros, atol=1, rtol=1e-7) + assert np.allclose(wrench_composer.local_force_b.numpy(), zeros, atol=1, rtol=1e-7) + assert np.allclose(wrench_composer.local_torque_b.numpy(), zeros, atol=1, rtol=1e-7) + assert np.allclose(wrench_composer.out_force_b.numpy(), zeros, atol=1, rtol=1e-7) + assert np.allclose(wrench_composer.out_torque_b.numpy(), zeros, atol=1, rtol=1e-7) # ============================================================================ @@ -404,13 +418,22 @@ def test_global_forces_with_rotation(device: str, num_envs: int, num_bodies: int forces_global = wp.from_numpy(forces_global_np, dtype=wp.vec3f, device=device) # Apply global forces - wrench_composer.add_forces_and_torques(forces=forces_global, is_global=True) + wrench_composer.add_forces_and_torques_index(forces=forces_global, is_global=True) # Compute expected local forces by rotating global forces by inverse quaternion expected_forces_local = quat_rotate_inv_np(link_quat_np, forces_global_np) + # Check raw global buffer has the global forces + global_force_np = wrench_composer.global_force_at_com_w.numpy() + assert np.allclose(global_force_np, forces_global_np, atol=1e-4, rtol=1e-5), ( + f"Global force buffer mismatch.\nExpected:\n{forces_global_np}\nGot:\n{global_force_np}" + ) + + # Compose to body frame before checking output + wrench_composer.compose_to_body_frame() + # Verify - composed_force_np = wrench_composer.composed_force.numpy() + composed_force_np = wrench_composer.out_force_b.numpy() assert np.allclose(composed_force_np, expected_forces_local, atol=1e-4, rtol=1e-5), ( f"Global force rotation failed.\nExpected:\n{expected_forces_local}\nGot:\n{composed_force_np}" ) @@ -437,13 +460,22 @@ def test_global_torques_with_rotation(device: str, num_envs: int, num_bodies: in torques_global = wp.from_numpy(torques_global_np, dtype=wp.vec3f, device=device) # Apply global torques - wrench_composer.add_forces_and_torques(torques=torques_global, is_global=True) + wrench_composer.add_forces_and_torques_index(torques=torques_global, is_global=True) # Compute expected local torques expected_torques_local = quat_rotate_inv_np(link_quat_np, torques_global_np) + # Check raw global buffer has the global torques + global_torque_np = wrench_composer.global_torque_w.numpy() + assert np.allclose(global_torque_np, torques_global_np, atol=1e-4, rtol=1e-5), ( + f"Global torque buffer mismatch.\nExpected:\n{torques_global_np}\nGot:\n{global_torque_np}" + ) + + # Compose to body frame before checking output + wrench_composer.compose_to_body_frame() + # Verify - composed_torque_np = wrench_composer.composed_torque.numpy() + composed_torque_np = wrench_composer.out_torque_b.numpy() assert np.allclose(composed_torque_np, expected_torques_local, atol=1e-4, rtol=1e-5), ( f"Global torque rotation failed.\nExpected:\n{expected_torques_local}\nGot:\n{composed_torque_np}" ) @@ -474,32 +506,40 @@ def test_global_forces_at_global_position(device: str, num_envs: int, num_bodies positions_global = wp.from_numpy(positions_global_np, dtype=wp.vec3f, device=device) # Apply global forces at global positions - wrench_composer.add_forces_and_torques(forces=forces_global, positions=positions_global, is_global=True) + wrench_composer.add_forces_and_torques_index(forces=forces_global, positions=positions_global, is_global=True) # Compute expected results: # 1. Force in local frame = quat_rotate_inv(link_quat, global_force) expected_forces_local = quat_rotate_inv_np(link_quat_np, forces_global_np) - # 2. Position offset in local frame = global_position - link_position (then used for torque) + # 2. Torque about CoM in world frame = cross(P_global - link_pos, F_global) + # Then rotate to body frame position_offset_global = positions_global_np - link_pos_np - - # 3. Torque = skew(position_offset_global) @ force_global, then rotate to local expected_torques_local = np.zeros((num_envs, num_bodies, 3), dtype=np.float32) for i in range(num_envs): for j in range(num_bodies): - pos_offset = position_offset_global[i, j] # global frame offset - force_local = expected_forces_local[i, j] # local frame force - # skew(pos_offset) @ force_local - expected_torques_local[i, j] = np.cross(pos_offset, force_local) + torque_w = np.cross(position_offset_global[i, j], forces_global_np[i, j]) + expected_torques_local[i, j] = quat_rotate_inv_np( + link_quat_np[i : i + 1, j : j + 1], torque_w.reshape(1, 1, 3) + )[0, 0] + + # Check raw global force buffer has the global forces + global_force_np = wrench_composer.global_force_w.numpy() + assert np.allclose(global_force_np, forces_global_np, atol=1e-4, rtol=1e-5), ( + f"Global force buffer mismatch.\nExpected:\n{forces_global_np}\nGot:\n{global_force_np}" + ) + + # Compose to body frame before checking output + wrench_composer.compose_to_body_frame() # Verify forces - composed_force_np = wrench_composer.composed_force.numpy() + composed_force_np = wrench_composer.out_force_b.numpy() assert np.allclose(composed_force_np, expected_forces_local, atol=1e-3, rtol=1e-4), ( f"Global force at position failed.\nExpected forces:\n{expected_forces_local}\nGot:\n{composed_force_np}" ) # Verify torques - composed_torque_np = wrench_composer.composed_torque.numpy() + composed_torque_np = wrench_composer.out_torque_b.numpy() assert np.allclose(composed_torque_np, expected_torques_local, atol=1e-3, rtol=1e-4), ( f"Global force at position failed.\nExpected torques:\n{expected_torques_local}\nGot:\n{composed_torque_np}" ) @@ -525,20 +565,24 @@ def test_local_vs_global_identity_quaternion(device: str): torques = wp.from_numpy(torques_np, dtype=wp.vec3f, device=device) # Apply as local - wrench_composer_local.add_forces_and_torques(forces=forces, torques=torques, is_global=False) + wrench_composer_local.add_forces_and_torques_index(forces=forces, torques=torques, is_global=False) # Apply as global (should be same with identity quaternion) - wrench_composer_global.add_forces_and_torques(forces=forces, torques=torques, is_global=True) + wrench_composer_global.add_forces_and_torques_index(forces=forces, torques=torques, is_global=True) + + # Compose to body frame before checking output + wrench_composer_local.compose_to_body_frame() + wrench_composer_global.compose_to_body_frame() # Results should be identical assert np.allclose( - wrench_composer_local.composed_force.numpy(), - wrench_composer_global.composed_force.numpy(), + wrench_composer_local.out_force_b.numpy(), + wrench_composer_global.out_force_b.numpy(), atol=1e-6, ) assert np.allclose( - wrench_composer_local.composed_torque.numpy(), - wrench_composer_global.composed_torque.numpy(), + wrench_composer_local.out_torque_b.numpy(), + wrench_composer_global.out_torque_b.numpy(), atol=1e-6, ) @@ -561,13 +605,16 @@ def test_90_degree_rotation_global_force(device: str): force_global = np.array([[[1.0, 0.0, 0.0]]], dtype=np.float32) force_wp = wp.from_numpy(force_global, dtype=wp.vec3f, device=device) - wrench_composer.add_forces_and_torques(forces=force_wp, is_global=True) + wrench_composer.add_forces_and_torques_index(forces=force_wp, is_global=True) # Expected: After inverse rotation (rotate by -90° around Z), X becomes -Y # Actually, inverse rotation of +90° around Z applied to (1,0,0) gives (0,-1,0) expected_force_local = np.array([[[0.0, -1.0, 0.0]]], dtype=np.float32) - composed_force_np = wrench_composer.composed_force.numpy() + # Compose to body frame before checking output + wrench_composer.compose_to_body_frame() + + composed_force_np = wrench_composer.out_force_b.numpy() assert np.allclose(composed_force_np, expected_force_local, atol=1e-5), ( f"90-degree rotation test failed.\nExpected:\n{expected_force_local}\nGot:\n{composed_force_np}" ) @@ -594,16 +641,25 @@ def test_composition_mixed_local_and_global(device: str): forces_global = wp.from_numpy(forces_global_np, dtype=wp.vec3f, device=device) # Add local forces first - wrench_composer.add_forces_and_torques(forces=forces_local, is_global=False) + wrench_composer.add_forces_and_torques_index(forces=forces_local, is_global=False) # Add global forces - wrench_composer.add_forces_and_torques(forces=forces_global, is_global=True) + wrench_composer.add_forces_and_torques_index(forces=forces_global, is_global=True) # Expected: local forces stay as-is, global forces get rotated, then sum global_forces_in_local = quat_rotate_inv_np(link_quat_np, forces_global_np) expected_total = forces_local_np + global_forces_in_local - composed_force_np = wrench_composer.composed_force.numpy() + # Check raw buffer properties + local_force_np = wrench_composer.local_force_b.numpy() + assert np.allclose(local_force_np, forces_local_np, atol=1e-4, rtol=1e-5) + global_force_at_com_np = wrench_composer.global_force_at_com_w.numpy() + assert np.allclose(global_force_at_com_np, forces_global_np, atol=1e-4, rtol=1e-5) + + # Compose to body frame before checking output + wrench_composer.compose_to_body_frame() + + composed_force_np = wrench_composer.out_force_b.numpy() assert np.allclose(composed_force_np, expected_total, atol=1e-4, rtol=1e-5), ( f"Mixed local/global composition failed.\nExpected:\n{expected_total}\nGot:\n{composed_force_np}" ) @@ -633,15 +689,22 @@ def test_local_forces_at_local_position(device: str, num_envs: int, num_bodies: positions_local = wp.from_numpy(positions_local_np, dtype=wp.vec3f, device=device) # Apply local forces at local positions - wrench_composer.add_forces_and_torques(forces=forces_local, positions=positions_local, is_global=False) + wrench_composer.add_forces_and_torques_index(forces=forces_local, positions=positions_local, is_global=False) # Expected: forces stay as-is, torque = cross(position, force) expected_forces = forces_local_np expected_torques = np.cross(positions_local_np, forces_local_np) + # Check raw local buffer + local_force_np = wrench_composer.local_force_b.numpy() + assert np.allclose(local_force_np, expected_forces, atol=1e-4, rtol=1e-5) + + # Compose to body frame before checking output + wrench_composer.compose_to_body_frame() + # Verify - composed_force_np = wrench_composer.composed_force.numpy() - composed_torque_np = wrench_composer.composed_torque.numpy() + composed_force_np = wrench_composer.out_force_b.numpy() + composed_torque_np = wrench_composer.out_torque_b.numpy() assert np.allclose(composed_force_np, expected_forces, atol=1e-4, rtol=1e-5) assert np.allclose(composed_torque_np, expected_torques, atol=1e-4, rtol=1e-5) @@ -670,14 +733,972 @@ def test_global_force_at_link_origin_no_torque(device: str): positions_at_link = wp.from_numpy(link_pos_np, dtype=wp.vec3f, device=device) # Apply global forces at link origin - wrench_composer.add_forces_and_torques(forces=forces_global, positions=positions_at_link, is_global=True) + wrench_composer.add_forces_and_torques_index(forces=forces_global, positions=positions_at_link, is_global=True) # Expected: force rotated to local, torque = 0 (since position offset is zero) expected_forces = quat_rotate_inv_np(link_quat_np, forces_global_np) expected_torques = np.zeros((num_envs, num_bodies, 3), dtype=np.float32) - composed_force_np = wrench_composer.composed_force.numpy() - composed_torque_np = wrench_composer.composed_torque.numpy() + # Check raw global force buffer + global_force_np = wrench_composer.global_force_w.numpy() + assert np.allclose(global_force_np, forces_global_np, atol=1e-4, rtol=1e-5) + + # Compose to body frame before checking output + wrench_composer.compose_to_body_frame() + + composed_force_np = wrench_composer.out_force_b.numpy() + composed_torque_np = wrench_composer.out_torque_b.numpy() assert np.allclose(composed_force_np, expected_forces, atol=1e-4, rtol=1e-5) assert np.allclose(composed_torque_np, expected_torques, atol=1e-4, rtol=1e-5) + + +# ============================================================================ +# add_raw_buffers_from Tests +# ============================================================================ + + +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.parametrize("num_envs", [1, 10, 100]) +@pytest.mark.parametrize("num_bodies", [1, 3, 5]) +def test_add_raw_buffers_from(device: str, num_envs: int, num_bodies: int): + """Test that add_raw_buffers_from merges all five input buffers correctly.""" + rng = np.random.default_rng(seed=20) + + # Create two composers with random link poses + link_pos_np = rng.uniform(-10.0, 10.0, (num_envs, num_bodies, 3)).astype(np.float32) + link_quat_np = random_unit_quaternion_np(rng, (num_envs, num_bodies)) + link_pos_torch = torch.from_numpy(link_pos_np) + link_quat_torch = torch.from_numpy(link_quat_np) + + mock_a = create_mock_asset(num_envs, num_bodies, device, link_pos=link_pos_torch, link_quat=link_quat_torch) + mock_b = create_mock_asset(num_envs, num_bodies, device, link_pos=link_pos_torch, link_quat=link_quat_torch) + + composer_a = WrenchComposer(mock_a) + composer_b = WrenchComposer(mock_b) + + # Populate composer_a with local forces at positions + forces_local_a_np = rng.uniform(-50.0, 50.0, (num_envs, num_bodies, 3)).astype(np.float32) + positions_local_a_np = rng.uniform(-5.0, 5.0, (num_envs, num_bodies, 3)).astype(np.float32) + composer_a.add_forces_and_torques_index( + forces=wp.from_numpy(forces_local_a_np, dtype=wp.vec3f, device=device), + positions=wp.from_numpy(positions_local_a_np, dtype=wp.vec3f, device=device), + is_global=False, + ) + + # Populate composer_b with global forces at global positions + forces_global_b_np = rng.uniform(-50.0, 50.0, (num_envs, num_bodies, 3)).astype(np.float32) + positions_global_b_np = rng.uniform(-5.0, 5.0, (num_envs, num_bodies, 3)).astype(np.float32) + composer_b.add_forces_and_torques_index( + forces=wp.from_numpy(forces_global_b_np, dtype=wp.vec3f, device=device), + positions=wp.from_numpy(positions_global_b_np, dtype=wp.vec3f, device=device), + is_global=True, + ) + + # Merge b into a + composer_a.add_raw_buffers_from(composer_b) + + # Build a reference composer that receives both calls directly + mock_ref = create_mock_asset(num_envs, num_bodies, device, link_pos=link_pos_torch, link_quat=link_quat_torch) + composer_ref = WrenchComposer(mock_ref) + composer_ref.add_forces_and_torques_index( + forces=wp.from_numpy(forces_local_a_np, dtype=wp.vec3f, device=device), + positions=wp.from_numpy(positions_local_a_np, dtype=wp.vec3f, device=device), + is_global=False, + ) + composer_ref.add_forces_and_torques_index( + forces=wp.from_numpy(forces_global_b_np, dtype=wp.vec3f, device=device), + positions=wp.from_numpy(positions_global_b_np, dtype=wp.vec3f, device=device), + is_global=True, + ) + + # Compose both and compare + composer_a.compose_to_body_frame() + composer_ref.compose_to_body_frame() + + assert np.allclose(composer_a.out_force_b.numpy(), composer_ref.out_force_b.numpy(), atol=1e-4, rtol=1e-5), ( + "add_raw_buffers_from force mismatch vs direct accumulation" + ) + assert np.allclose(composer_a.out_torque_b.numpy(), composer_ref.out_torque_b.numpy(), atol=1e-4, rtol=1e-5), ( + "add_raw_buffers_from torque mismatch vs direct accumulation" + ) + + +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +def test_add_raw_buffers_from_inactive_is_noop(device: str): + """Test that add_raw_buffers_from is a no-op when the source composer is inactive.""" + num_envs, num_bodies = 4, 2 + rng = np.random.default_rng(seed=21) + + mock_a = create_mock_asset(num_envs, num_bodies, device) + mock_b = create_mock_asset(num_envs, num_bodies, device) + composer_a = WrenchComposer(mock_a) + composer_b = WrenchComposer(mock_b) + + # Populate composer_a with some forces + forces_np = rng.uniform(-50.0, 50.0, (num_envs, num_bodies, 3)).astype(np.float32) + composer_a.add_forces_and_torques_index( + forces=wp.from_numpy(forces_np, dtype=wp.vec3f, device=device), + ) + + # composer_b is inactive (never written to) + assert not composer_b.active + + # Snapshot composer_a's local buffer before merge + local_force_before = composer_a.local_force_b.numpy().copy() + + # Merge inactive composer_b into composer_a -- should be a no-op + composer_a.add_raw_buffers_from(composer_b) + + assert np.allclose(composer_a.local_force_b.numpy(), local_force_before, atol=1e-7) + + +# ============================================================================ +# Mask-based API Tests +# ============================================================================ + + +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.parametrize("num_envs", [1, 10, 100]) +@pytest.mark.parametrize("num_bodies", [1, 3, 5]) +def test_add_forces_mask(device: str, num_envs: int, num_bodies: int): + """Test that add_forces_and_torques_mask produces the same result as the index variant.""" + rng = np.random.default_rng(seed=30) + + for _ in range(5): + # Random subset selection + env_select = rng.choice([True, False], size=num_envs, replace=True) + body_select = rng.choice([True, False], size=num_bodies, replace=True) + # Ensure at least one env and body are selected + env_select[0] = True + body_select[0] = True + + env_ids_np = np.where(env_select)[0].astype(np.int32) + body_ids_np = np.where(body_select)[0].astype(np.int32) + env_mask_np = env_select.astype(np.bool_) + body_mask_np = body_select.astype(np.bool_) + + # Random forces for the full grid (mask variant takes full-sized arrays) + forces_full_np = rng.uniform(-100.0, 100.0, (num_envs, num_bodies, 3)).astype(np.float32) + + # Index-based composer + mock_idx = create_mock_asset(num_envs, num_bodies, device) + composer_idx = WrenchComposer(mock_idx) + # Extract the subset for index API + forces_subset_np = forces_full_np[env_ids_np[:, None], body_ids_np[None, :], :] + composer_idx.add_forces_and_torques_index( + forces=wp.from_numpy(forces_subset_np, dtype=wp.vec3f, device=device), + env_ids=wp.from_numpy(env_ids_np, dtype=wp.int32, device=device), + body_ids=wp.from_numpy(body_ids_np, dtype=wp.int32, device=device), + ) + + # Mask-based composer + mock_mask = create_mock_asset(num_envs, num_bodies, device) + composer_mask = WrenchComposer(mock_mask) + composer_mask.add_forces_and_torques_mask( + forces=wp.from_numpy(forces_full_np, dtype=wp.vec3f, device=device), + env_mask=wp.from_numpy(env_mask_np, dtype=wp.bool, device=device), + body_mask=wp.from_numpy(body_mask_np, dtype=wp.bool, device=device), + ) + + # Compose both + composer_idx.compose_to_body_frame() + composer_mask.compose_to_body_frame() + + assert np.allclose(composer_idx.out_force_b.numpy(), composer_mask.out_force_b.numpy(), atol=1e-4, rtol=1e-5), ( + f"Mask vs index force mismatch (envs={num_envs}, bodies={num_bodies})" + ) + assert np.allclose( + composer_idx.out_torque_b.numpy(), composer_mask.out_torque_b.numpy(), atol=1e-4, rtol=1e-5 + ), f"Mask vs index torque mismatch (envs={num_envs}, bodies={num_bodies})" + + +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +@pytest.mark.parametrize("num_envs", [1, 10, 100]) +@pytest.mark.parametrize("num_bodies", [1, 3, 5]) +def test_add_forces_mask_global(device: str, num_envs: int, num_bodies: int): + """Test mask-based API with global forces and positions.""" + rng = np.random.default_rng(seed=31) + + # Random link poses + link_pos_np = rng.uniform(-10.0, 10.0, (num_envs, num_bodies, 3)).astype(np.float32) + link_quat_np = random_unit_quaternion_np(rng, (num_envs, num_bodies)) + link_pos_torch = torch.from_numpy(link_pos_np) + link_quat_torch = torch.from_numpy(link_quat_np) + + # Select all envs and bodies to keep comparison simple + forces_np = rng.uniform(-100.0, 100.0, (num_envs, num_bodies, 3)).astype(np.float32) + positions_np = rng.uniform(-10.0, 10.0, (num_envs, num_bodies, 3)).astype(np.float32) + + # Index-based + mock_idx = create_mock_asset(num_envs, num_bodies, device, link_pos=link_pos_torch, link_quat=link_quat_torch) + composer_idx = WrenchComposer(mock_idx) + composer_idx.add_forces_and_torques_index( + forces=wp.from_numpy(forces_np, dtype=wp.vec3f, device=device), + positions=wp.from_numpy(positions_np, dtype=wp.vec3f, device=device), + is_global=True, + ) + + # Mask-based (all-True masks) + mock_mask = create_mock_asset(num_envs, num_bodies, device, link_pos=link_pos_torch, link_quat=link_quat_torch) + composer_mask = WrenchComposer(mock_mask) + env_mask = wp.from_numpy(np.ones(num_envs, dtype=np.bool_), dtype=wp.bool, device=device) + body_mask = wp.from_numpy(np.ones(num_bodies, dtype=np.bool_), dtype=wp.bool, device=device) + composer_mask.add_forces_and_torques_mask( + forces=wp.from_numpy(forces_np, dtype=wp.vec3f, device=device), + positions=wp.from_numpy(positions_np, dtype=wp.vec3f, device=device), + env_mask=env_mask, + body_mask=body_mask, + is_global=True, + ) + + composer_idx.compose_to_body_frame() + composer_mask.compose_to_body_frame() + + assert np.allclose(composer_idx.out_force_b.numpy(), composer_mask.out_force_b.numpy(), atol=1e-4, rtol=1e-5), ( + "Mask vs index global force mismatch" + ) + assert np.allclose(composer_idx.out_torque_b.numpy(), composer_mask.out_torque_b.numpy(), atol=1e-4, rtol=1e-5), ( + "Mask vs index global torque mismatch" + ) + + +# ============================================================================ +# set_forces_and_torques_index Tests +# ============================================================================ + + +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +def test_set_forces_overwrites_previous_add(device: str): + """Test that set_forces_and_torques_index clears previously accumulated values.""" + num_envs, num_bodies = 4, 2 + rng = np.random.default_rng(seed=40) + + mock_asset = create_mock_asset(num_envs, num_bodies, device) + composer = WrenchComposer(mock_asset) + + # First accumulate some forces via add + forces_a_np = rng.uniform(-100.0, 100.0, (num_envs, num_bodies, 3)).astype(np.float32) + composer.add_forces_and_torques_index( + forces=wp.from_numpy(forces_a_np, dtype=wp.vec3f, device=device), + ) + + # Now set new forces -- should replace, not accumulate + forces_b_np = rng.uniform(-100.0, 100.0, (num_envs, num_bodies, 3)).astype(np.float32) + composer.set_forces_and_torques_index( + forces=wp.from_numpy(forces_b_np, dtype=wp.vec3f, device=device), + ) + + composer.compose_to_body_frame() + + # Output should match forces_b only (forces_a should be gone) + assert np.allclose(composer.out_force_b.numpy(), forces_b_np, atol=1e-4, rtol=1e-5), ( + "set_forces did not clear previous add" + ) + + +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +def test_set_forces_clears_targeted_envs_only(device: str): + """Test that set_forces_and_torques_index clears only the targeted environments.""" + num_envs, num_bodies = 4, 3 + rng = np.random.default_rng(seed=41) + + mock_asset = create_mock_asset(num_envs, num_bodies, device) + composer = WrenchComposer(mock_asset) + + # Add global forces at positions (populates global_force_w and global_torque_w) + forces_global_np = rng.uniform(-50.0, 50.0, (num_envs, num_bodies, 3)).astype(np.float32) + positions_np = rng.uniform(-5.0, 5.0, (num_envs, num_bodies, 3)).astype(np.float32) + composer.add_forces_and_torques_index( + forces=wp.from_numpy(forces_global_np, dtype=wp.vec3f, device=device), + positions=wp.from_numpy(positions_np, dtype=wp.vec3f, device=device), + is_global=True, + ) + + # Also add local torques (populates local_torque_b) + torques_local_np = rng.uniform(-50.0, 50.0, (num_envs, num_bodies, 3)).astype(np.float32) + composer.add_forces_and_torques_index( + torques=wp.from_numpy(torques_local_np, dtype=wp.vec3f, device=device), + is_global=False, + ) + + # Now set local forces for envs [0, 2] -- should clear only envs 0, 2 + env_ids_np = np.array([0, 2], dtype=np.int32) + kept_env_ids = np.array([1, 3], dtype=np.int32) + forces_new_np = rng.uniform(-50.0, 50.0, (2, num_bodies, 3)).astype(np.float32) + composer.set_forces_and_torques_index( + forces=wp.from_numpy(forces_new_np, dtype=wp.vec3f, device=device), + env_ids=wp.from_numpy(env_ids_np, dtype=wp.int32, device=device), + is_global=False, + ) + + zeros = np.zeros((num_bodies, 3), dtype=np.float32) + + # Targeted envs [0, 2]: all buffers cleared, then local_force_b written + for eid in env_ids_np: + assert np.allclose(composer.global_force_w.numpy()[eid], zeros, atol=1e-7), ( + f"global_force_w not cleared for targeted env {eid}" + ) + assert np.allclose(composer.global_torque_w.numpy()[eid], zeros, atol=1e-7), ( + f"global_torque_w not cleared for targeted env {eid}" + ) + assert np.allclose(composer.local_torque_b.numpy()[eid], zeros, atol=1e-7), ( + f"local_torque_b not cleared for targeted env {eid}" + ) + + # Non-targeted envs [1, 3]: should retain original values + for eid in kept_env_ids: + assert np.allclose(composer.global_force_w.numpy()[eid], forces_global_np[eid], atol=1e-4, rtol=1e-5), ( + f"global_force_w changed for non-targeted env {eid}" + ) + assert np.allclose(composer.local_torque_b.numpy()[eid], torques_local_np[eid], atol=1e-4, rtol=1e-5), ( + f"local_torque_b changed for non-targeted env {eid}" + ) + + # local_force_b should have new values at env_ids [0, 2], zeros at [1, 3] + expected_local_force = np.zeros((num_envs, num_bodies, 3), dtype=np.float32) + expected_local_force[env_ids_np] = forces_new_np + assert np.allclose(composer.local_force_b.numpy(), expected_local_force, atol=1e-4, rtol=1e-5), ( + "local_force_b has wrong values after set" + ) + + +# ============================================================================ +# Partial Reset Tests +# ============================================================================ + + +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +def test_partial_reset_zeros_only_specified_envs(device: str): + """Test that partial reset zeros only the specified environments and leaves others intact.""" + num_envs, num_bodies = 8, 3 + rng = np.random.default_rng(seed=50) + + mock_asset = create_mock_asset(num_envs, num_bodies, device) + composer = WrenchComposer(mock_asset) + + # Populate all envs with local forces + forces_np = rng.uniform(-100.0, 100.0, (num_envs, num_bodies, 3)).astype(np.float32) + composer.add_forces_and_torques_index( + forces=wp.from_numpy(forces_np, dtype=wp.vec3f, device=device), + ) + + # Also add global forces to populate more buffers + forces_global_np = rng.uniform(-100.0, 100.0, (num_envs, num_bodies, 3)).astype(np.float32) + composer.add_forces_and_torques_index( + forces=wp.from_numpy(forces_global_np, dtype=wp.vec3f, device=device), + is_global=True, + ) + + # Partial reset: only envs [1, 3, 5] + reset_env_ids = np.array([1, 3, 5], dtype=np.int32) + kept_env_ids = np.array([0, 2, 4, 6, 7], dtype=np.int32) + composer.reset(env_ids=wp.from_numpy(reset_env_ids, dtype=wp.int32, device=device)) + + # Reset envs should be zeroed across all input buffers + zeros = np.zeros((num_bodies, 3), dtype=np.float32) + local_force = composer.local_force_b.numpy() + global_force_at_com = composer.global_force_at_com_w.numpy() + for eid in reset_env_ids: + assert np.allclose(local_force[eid], zeros, atol=1e-7), f"local_force_b not zeroed for env {eid}" + assert np.allclose(global_force_at_com[eid], zeros, atol=1e-7), ( + f"global_force_at_com_w not zeroed for env {eid}" + ) + + # Kept envs should retain their values + for eid in kept_env_ids: + assert np.allclose(local_force[eid], forces_np[eid], atol=1e-4, rtol=1e-5), ( + f"local_force_b changed for non-reset env {eid}" + ) + assert np.allclose(global_force_at_com[eid], forces_global_np[eid], atol=1e-4, rtol=1e-5), ( + f"global_force_at_com_w changed for non-reset env {eid}" + ) + + # Flags: _active should still be True, _dirty should be True + assert composer.active + assert composer._dirty + + +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +def test_full_reset_clears_active_flag(device: str): + """Test that full reset (no args) clears the _active flag.""" + num_envs, num_bodies = 4, 2 + + mock_asset = create_mock_asset(num_envs, num_bodies, device) + composer = WrenchComposer(mock_asset) + + forces_np = np.ones((num_envs, num_bodies, 3), dtype=np.float32) + composer.add_forces_and_torques_index( + forces=wp.from_numpy(forces_np, dtype=wp.vec3f, device=device), + ) + assert composer.active + + composer.reset() + assert not composer.active + assert not composer._dirty + + +# ============================================================================ +# Deprecated API Backward-Compatibility Tests +# ============================================================================ + + +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +def test_composed_force_emits_deprecation_warning(device: str): + """Test that accessing composed_force emits a DeprecationWarning.""" + num_envs, num_bodies = 2, 1 + + mock_asset = create_mock_asset(num_envs, num_bodies, device) + composer = WrenchComposer(mock_asset) + + forces_np = np.array([[[1.0, 2.0, 3.0]], [[4.0, 5.0, 6.0]]], dtype=np.float32) + composer.add_forces_and_torques_index( + forces=wp.from_numpy(forces_np, dtype=wp.vec3f, device=device), + ) + + with pytest.warns(DeprecationWarning, match="composed_force.*is deprecated"): + result = composer.composed_force + + # Should return the same data as out_force_b + assert np.allclose(result.numpy(), composer.out_force_b.numpy(), atol=1e-7) + + +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +def test_composed_torque_emits_deprecation_warning(device: str): + """Test that accessing composed_torque emits a DeprecationWarning.""" + num_envs, num_bodies = 2, 1 + + mock_asset = create_mock_asset(num_envs, num_bodies, device) + composer = WrenchComposer(mock_asset) + + torques_np = np.array([[[1.0, 2.0, 3.0]], [[4.0, 5.0, 6.0]]], dtype=np.float32) + composer.add_forces_and_torques_index( + torques=wp.from_numpy(torques_np, dtype=wp.vec3f, device=device), + ) + + with pytest.warns(DeprecationWarning, match="composed_torque.*is deprecated"): + result = composer.composed_torque + + assert np.allclose(result.numpy(), composer.out_torque_b.numpy(), atol=1e-7) + + +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +def test_deprecated_add_forces_and_torques_emits_warning(device: str): + """Test that the deprecated add_forces_and_torques wrapper emits a warning and works.""" + num_envs, num_bodies = 4, 2 + rng = np.random.default_rng(seed=52) + + mock_asset = create_mock_asset(num_envs, num_bodies, device) + composer = WrenchComposer(mock_asset) + + forces_np = rng.uniform(-50.0, 50.0, (num_envs, num_bodies, 3)).astype(np.float32) + + with pytest.warns(DeprecationWarning, match="add_forces_and_torques.*is deprecated"): + composer.add_forces_and_torques( + forces=wp.from_numpy(forces_np, dtype=wp.vec3f, device=device), + ) + + composer.compose_to_body_frame() + assert np.allclose(composer.out_force_b.numpy(), forces_np, atol=1e-4, rtol=1e-5) + + +# ============================================================================ +# set_forces_and_torques_mask Tests +# ============================================================================ + + +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +def test_set_forces_mask_overwrites_previous_add(device: str): + """Test that set_forces_and_torques_mask clears previously accumulated values.""" + num_envs, num_bodies = 4, 2 + rng = np.random.default_rng(seed=60) + + mock_asset = create_mock_asset(num_envs, num_bodies, device) + composer = WrenchComposer(mock_asset) + + # Accumulate some forces via add + forces_a_np = rng.uniform(-100.0, 100.0, (num_envs, num_bodies, 3)).astype(np.float32) + composer.add_forces_and_torques_index( + forces=wp.from_numpy(forces_a_np, dtype=wp.vec3f, device=device), + ) + + # Now set new forces via mask -- should replace, not accumulate + forces_b_np = rng.uniform(-100.0, 100.0, (num_envs, num_bodies, 3)).astype(np.float32) + composer.set_forces_and_torques_mask( + forces=wp.from_numpy(forces_b_np, dtype=wp.vec3f, device=device), + ) + + composer.compose_to_body_frame() + + # Output should match forces_b only (forces_a should be gone) + assert np.allclose(composer.out_force_b.numpy(), forces_b_np, atol=1e-4, rtol=1e-5), ( + "set_forces_and_torques_mask did not clear previous add" + ) + + +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +def test_set_forces_mask_clears_targeted_envs_only(device: str): + """Test that set_forces_and_torques_mask clears only the masked environments.""" + num_envs, num_bodies = 4, 3 + rng = np.random.default_rng(seed=61) + + mock_asset = create_mock_asset(num_envs, num_bodies, device) + composer = WrenchComposer(mock_asset) + + # Populate global buffers for all envs + forces_global_np = rng.uniform(-50.0, 50.0, (num_envs, num_bodies, 3)).astype(np.float32) + positions_np = rng.uniform(-5.0, 5.0, (num_envs, num_bodies, 3)).astype(np.float32) + composer.add_forces_and_torques_index( + forces=wp.from_numpy(forces_global_np, dtype=wp.vec3f, device=device), + positions=wp.from_numpy(positions_np, dtype=wp.vec3f, device=device), + is_global=True, + ) + + # Also add local torques for all envs + torques_local_np = rng.uniform(-50.0, 50.0, (num_envs, num_bodies, 3)).astype(np.float32) + composer.add_forces_and_torques_index( + torques=wp.from_numpy(torques_local_np, dtype=wp.vec3f, device=device), + is_global=False, + ) + + # Set local forces via mask for envs [0, 2] -- should clear only masked envs + env_mask_np = np.array([True, False, True, False], dtype=np.bool_) + body_mask_np = np.array([True, True, False], dtype=np.bool_) + forces_new_np = rng.uniform(-50.0, 50.0, (num_envs, num_bodies, 3)).astype(np.float32) + composer.set_forces_and_torques_mask( + forces=wp.from_numpy(forces_new_np, dtype=wp.vec3f, device=device), + env_mask=wp.from_numpy(env_mask_np, dtype=wp.bool, device=device), + body_mask=wp.from_numpy(body_mask_np, dtype=wp.bool, device=device), + is_global=False, + ) + + zeros = np.zeros((num_bodies, 3), dtype=np.float32) + + # Masked envs [0, 2]: all buffers cleared by reset, then local_force_b written where body_mask is True + for eid in [0, 2]: + assert np.allclose(composer.global_force_w.numpy()[eid], zeros, atol=1e-7), ( + f"global_force_w not cleared for masked env {eid}" + ) + assert np.allclose(composer.global_torque_w.numpy()[eid], zeros, atol=1e-7), ( + f"global_torque_w not cleared for masked env {eid}" + ) + assert np.allclose(composer.local_torque_b.numpy()[eid], zeros, atol=1e-7), ( + f"local_torque_b not cleared for masked env {eid}" + ) + + # Non-masked envs [1, 3]: should retain original values + for eid in [1, 3]: + assert np.allclose(composer.global_force_w.numpy()[eid], forces_global_np[eid], atol=1e-4, rtol=1e-5), ( + f"global_force_w changed for non-masked env {eid}" + ) + assert np.allclose(composer.local_torque_b.numpy()[eid], torques_local_np[eid], atol=1e-4, rtol=1e-5), ( + f"local_torque_b changed for non-masked env {eid}" + ) + + # local_force_b should have new values where both masks are True, zeros for masked envs otherwise + expected_local_force = np.zeros((num_envs, num_bodies, 3), dtype=np.float32) + for e in range(num_envs): + for b in range(num_bodies): + if env_mask_np[e] and body_mask_np[b]: + expected_local_force[e, b] = forces_new_np[e, b] + assert np.allclose(composer.local_force_b.numpy(), expected_local_force, atol=1e-4, rtol=1e-5), ( + "local_force_b has wrong values after mask set" + ) + + +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +def test_set_forces_mask_matches_set_forces_index(device: str): + """Test that set_forces_and_torques_mask produces the same result as the index variant.""" + num_envs, num_bodies = 6, 3 + rng = np.random.default_rng(seed=62) + + # Random link poses + link_pos_np = rng.uniform(-10.0, 10.0, (num_envs, num_bodies, 3)).astype(np.float32) + link_quat_np = random_unit_quaternion_np(rng, (num_envs, num_bodies)) + link_pos_torch = torch.from_numpy(link_pos_np) + link_quat_torch = torch.from_numpy(link_quat_np) + + # Use all envs/bodies to compare + forces_np = rng.uniform(-100.0, 100.0, (num_envs, num_bodies, 3)).astype(np.float32) + positions_np = rng.uniform(-10.0, 10.0, (num_envs, num_bodies, 3)).astype(np.float32) + + # Index-based + mock_idx = create_mock_asset(num_envs, num_bodies, device, link_pos=link_pos_torch, link_quat=link_quat_torch) + composer_idx = WrenchComposer(mock_idx) + composer_idx.set_forces_and_torques_index( + forces=wp.from_numpy(forces_np, dtype=wp.vec3f, device=device), + positions=wp.from_numpy(positions_np, dtype=wp.vec3f, device=device), + is_global=True, + ) + + # Mask-based (all-True) + mock_mask = create_mock_asset(num_envs, num_bodies, device, link_pos=link_pos_torch, link_quat=link_quat_torch) + composer_mask = WrenchComposer(mock_mask) + composer_mask.set_forces_and_torques_mask( + forces=wp.from_numpy(forces_np, dtype=wp.vec3f, device=device), + positions=wp.from_numpy(positions_np, dtype=wp.vec3f, device=device), + is_global=True, + ) + + composer_idx.compose_to_body_frame() + composer_mask.compose_to_body_frame() + + assert np.allclose(composer_idx.out_force_b.numpy(), composer_mask.out_force_b.numpy(), atol=1e-4, rtol=1e-5), ( + "set mask vs index force mismatch" + ) + assert np.allclose(composer_idx.out_torque_b.numpy(), composer_mask.out_torque_b.numpy(), atol=1e-4, rtol=1e-5), ( + "set mask vs index torque mismatch" + ) + + +# ============================================================================ +# Lazy Composition (_ensure_composed) Tests +# ============================================================================ + + +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +def test_out_force_b_triggers_lazy_composition(device: str): + """Test that accessing out_force_b without explicit compose_to_body_frame still returns correct results.""" + num_envs, num_bodies = 4, 2 + rng = np.random.default_rng(seed=70) + + link_quat_np = random_unit_quaternion_np(rng, (num_envs, num_bodies)) + link_quat_torch = torch.from_numpy(link_quat_np) + + mock_asset = create_mock_asset(num_envs, num_bodies, device, link_quat=link_quat_torch) + composer = WrenchComposer(mock_asset) + + forces_global_np = rng.uniform(-100.0, 100.0, (num_envs, num_bodies, 3)).astype(np.float32) + composer.add_forces_and_torques_index( + forces=wp.from_numpy(forces_global_np, dtype=wp.vec3f, device=device), + is_global=True, + ) + + # Do NOT call compose_to_body_frame -- rely on lazy composition + expected_forces_local = quat_rotate_inv_np(link_quat_np, forces_global_np) + composed_force_np = composer.out_force_b.numpy() + + assert np.allclose(composed_force_np, expected_forces_local, atol=1e-4, rtol=1e-5), ( + "Lazy composition via out_force_b failed" + ) + + +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +def test_out_torque_b_triggers_lazy_composition(device: str): + """Test that accessing out_torque_b without explicit compose_to_body_frame still returns correct results.""" + num_envs, num_bodies = 4, 2 + rng = np.random.default_rng(seed=71) + + link_quat_np = random_unit_quaternion_np(rng, (num_envs, num_bodies)) + link_quat_torch = torch.from_numpy(link_quat_np) + + mock_asset = create_mock_asset(num_envs, num_bodies, device, link_quat=link_quat_torch) + composer = WrenchComposer(mock_asset) + + torques_global_np = rng.uniform(-100.0, 100.0, (num_envs, num_bodies, 3)).astype(np.float32) + composer.add_forces_and_torques_index( + torques=wp.from_numpy(torques_global_np, dtype=wp.vec3f, device=device), + is_global=True, + ) + + # Do NOT call compose_to_body_frame -- rely on lazy composition + expected_torques_local = quat_rotate_inv_np(link_quat_np, torques_global_np) + composed_torque_np = composer.out_torque_b.numpy() + + assert np.allclose(composed_torque_np, expected_torques_local, atol=1e-4, rtol=1e-5), ( + "Lazy composition via out_torque_b failed" + ) + + +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +def test_lazy_composition_tracks_dirty_flag(device: str): + """Test that the dirty flag is correctly managed through add/compose/add cycles.""" + num_envs, num_bodies = 2, 1 + + mock_asset = create_mock_asset(num_envs, num_bodies, device) + composer = WrenchComposer(mock_asset) + + # Initially clean + assert not composer._dirty + + # After add, dirty + forces_np = np.ones((num_envs, num_bodies, 3), dtype=np.float32) + composer.add_forces_and_torques_index( + forces=wp.from_numpy(forces_np, dtype=wp.vec3f, device=device), + ) + assert composer._dirty + + # After accessing out_force_b, clean (lazy compose happened) + _ = composer.out_force_b + assert not composer._dirty + + # After another add, dirty again + composer.add_forces_and_torques_index( + forces=wp.from_numpy(forces_np, dtype=wp.vec3f, device=device), + ) + assert composer._dirty + + # Accessing out_torque_b also triggers composition + _ = composer.out_torque_b + assert not composer._dirty + + # Verify accumulated result (2x forces) + expected = 2.0 * forces_np + assert np.allclose(composer.out_force_b.numpy(), expected, atol=1e-4, rtol=1e-5) + + +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +def test_compose_is_idempotent(device: str): + """Calling compose_to_body_frame twice without intervening writes produces the same result.""" + rng = np.random.default_rng(seed=456) + num_envs, num_bodies = 4, 3 + + # Non-trivial link pose so the rotation path is exercised + link_pos_np = rng.uniform(-2, 2, (num_envs, num_bodies, 3)).astype(np.float32) + link_quat_np = rng.standard_normal((num_envs, num_bodies, 4)).astype(np.float32) + link_quat_np /= np.linalg.norm(link_quat_np, axis=-1, keepdims=True) + + mock_asset = create_mock_asset( + num_envs, + num_bodies, + device, + link_pos=torch.from_numpy(link_pos_np), + link_quat=torch.from_numpy(link_quat_np), + ) + composer = WrenchComposer(mock_asset) + + # Add global forces with positions (exercises cross-product torque path) + forces_np = rng.uniform(-5, 5, (num_envs, num_bodies, 3)).astype(np.float32) + positions_np = rng.uniform(-1, 1, (num_envs, num_bodies, 3)).astype(np.float32) + torques_np = rng.uniform(-3, 3, (num_envs, num_bodies, 3)).astype(np.float32) + + composer.add_forces_and_torques_index( + forces=wp.from_numpy(forces_np, dtype=wp.vec3f, device=device), + torques=wp.from_numpy(torques_np, dtype=wp.vec3f, device=device), + positions=wp.from_numpy(positions_np, dtype=wp.vec3f, device=device), + is_global=True, + ) + + # First compose + composer.compose_to_body_frame() + force_first = composer.out_force_b.numpy().copy() + torque_first = composer.out_torque_b.numpy().copy() + + # Second compose (no writes in between) + composer.compose_to_body_frame() + force_second = composer.out_force_b.numpy() + torque_second = composer.out_torque_b.numpy() + + np.testing.assert_array_equal(force_first, force_second) + np.testing.assert_array_equal(torque_first, torque_second) + + +# ============================================================================ +# CoM Offset from Link Origin Tests +# ============================================================================ + + +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +def test_global_force_with_com_offset(device: str): + """Test that torque correction uses CoM position, not link position, when they differ.""" + num_envs, num_bodies = 2, 1 + + # Link at origin, CoM offset by [1, 0, 0] + link_pos_np = np.zeros((num_envs, num_bodies, 3), dtype=np.float32) + link_quat_np = np.zeros((num_envs, num_bodies, 4), dtype=np.float32) + link_quat_np[..., 3] = 1.0 # identity quaternion (xyzw) + + com_pos_np = np.zeros((num_envs, num_bodies, 3), dtype=np.float32) + com_pos_np[..., 0] = 1.0 # CoM at [1, 0, 0] + + mock_asset = create_mock_asset( + num_envs, + num_bodies, + device, + link_pos=torch.from_numpy(link_pos_np), + link_quat=torch.from_numpy(link_quat_np), + ) + # Set CoM pose separately (pos=[1,0,0], quat=identity) + com_pose = torch.cat([torch.from_numpy(com_pos_np), torch.from_numpy(link_quat_np)], dim=-1) + mock_asset.data.set_body_com_pose_w(com_pose) + + composer = WrenchComposer(mock_asset) + + # Apply global force [0, 0, 10] at position [0, 0, 0] (world origin) + forces_np = np.zeros((num_envs, num_bodies, 3), dtype=np.float32) + forces_np[..., 2] = 10.0 + positions_np = np.zeros((num_envs, num_bodies, 3), dtype=np.float32) + + composer.add_forces_and_torques_index( + forces=wp.from_numpy(forces_np, dtype=wp.vec3f, device=device), + positions=wp.from_numpy(positions_np, dtype=wp.vec3f, device=device), + is_global=True, + ) + + composer.compose_to_body_frame() + + # With identity quaternion: + # torque_w = cross(P, F) - cross(com, F) = cross([0,0,0], [0,0,10]) - cross([1,0,0], [0,0,10]) + # = [0,0,0] - [0*10-0*0, 0*0-1*10, 1*0-0*0] = [0,0,0] - [0, -10, 0] = [0, 10, 0] + # In body frame (identity rotation): [0, 10, 0] + expected_torque = np.zeros((num_envs, num_bodies, 3), dtype=np.float32) + expected_torque[..., 1] = 10.0 + + assert np.allclose(composer.out_torque_b.numpy(), expected_torque, atol=1e-4, rtol=1e-5), ( + f"CoM offset torque correction failed.\nExpected:\n{expected_torque}\nGot:\n{composer.out_torque_b.numpy()}" + ) + + # Force should be unchanged (identity rotation) + assert np.allclose(composer.out_force_b.numpy(), forces_np, atol=1e-4, rtol=1e-5) + + +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +def test_global_force_at_com_no_torque_with_com_offset(device: str): + """Test that a global force at CoM position produces zero torque even with CoM offset.""" + num_envs, num_bodies = 2, 1 + + # Link at origin, CoM offset by [2, 3, 0] + link_pos_np = np.zeros((num_envs, num_bodies, 3), dtype=np.float32) + link_quat_np = np.zeros((num_envs, num_bodies, 4), dtype=np.float32) + link_quat_np[..., 3] = 1.0 + + com_pos_np = np.zeros((num_envs, num_bodies, 3), dtype=np.float32) + com_pos_np[..., 0] = 2.0 + com_pos_np[..., 1] = 3.0 + + mock_asset = create_mock_asset( + num_envs, + num_bodies, + device, + link_pos=torch.from_numpy(link_pos_np), + link_quat=torch.from_numpy(link_quat_np), + ) + com_pose = torch.cat([torch.from_numpy(com_pos_np), torch.from_numpy(link_quat_np)], dim=-1) + mock_asset.data.set_body_com_pose_w(com_pose) + + composer = WrenchComposer(mock_asset) + + # Apply global force at the CoM position + forces_np = np.zeros((num_envs, num_bodies, 3), dtype=np.float32) + forces_np[..., 2] = 50.0 + positions_np = com_pos_np.copy() + + composer.add_forces_and_torques_index( + forces=wp.from_numpy(forces_np, dtype=wp.vec3f, device=device), + positions=wp.from_numpy(positions_np, dtype=wp.vec3f, device=device), + is_global=True, + ) + + composer.compose_to_body_frame() + + # Torque = cross(com, F) - cross(com, F) = 0 + expected_torque = np.zeros((num_envs, num_bodies, 3), dtype=np.float32) + assert np.allclose(composer.out_torque_b.numpy(), expected_torque, atol=1e-4, rtol=1e-5), ( + "Force at CoM should produce zero torque regardless of CoM offset" + ) + + +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +def test_com_offset_with_rotation(device: str): + """Test torque correction with both CoM offset and non-identity rotation.""" + num_envs, num_bodies = 1, 1 + rng = np.random.default_rng(seed=73) + + # Random rotation + link_quat_np = random_unit_quaternion_np(rng, (num_envs, num_bodies)) + link_pos_np = rng.uniform(-5.0, 5.0, (num_envs, num_bodies, 3)).astype(np.float32) + + # CoM offset from link + com_offset_np = rng.uniform(0.5, 2.0, (num_envs, num_bodies, 3)).astype(np.float32) + com_pos_np = link_pos_np + com_offset_np # simple world-frame offset for test clarity + + mock_asset = create_mock_asset( + num_envs, + num_bodies, + device, + link_pos=torch.from_numpy(link_pos_np), + link_quat=torch.from_numpy(link_quat_np), + ) + com_pose = torch.cat([torch.from_numpy(com_pos_np), torch.from_numpy(link_quat_np)], dim=-1) + mock_asset.data.set_body_com_pose_w(com_pose) + + composer = WrenchComposer(mock_asset) + + # Apply global force at a random world position + forces_np = rng.uniform(-100.0, 100.0, (num_envs, num_bodies, 3)).astype(np.float32) + positions_np = rng.uniform(-10.0, 10.0, (num_envs, num_bodies, 3)).astype(np.float32) + + composer.add_forces_and_torques_index( + forces=wp.from_numpy(forces_np, dtype=wp.vec3f, device=device), + positions=wp.from_numpy(positions_np, dtype=wp.vec3f, device=device), + is_global=True, + ) + + composer.compose_to_body_frame() + + # Expected: torque_w = cross(P, F) - cross(com, F) = cross(P - com, F) + lever_arm = positions_np - com_pos_np + torque_w = np.cross(lever_arm, forces_np) + expected_torque_b = quat_rotate_inv_np(link_quat_np, torque_w) + expected_force_b = quat_rotate_inv_np(link_quat_np, forces_np) + + assert np.allclose(composer.out_force_b.numpy(), expected_force_b, atol=1e-3, rtol=1e-4), ( + "Force mismatch with CoM offset + rotation" + ) + assert np.allclose(composer.out_torque_b.numpy(), expected_torque_b, atol=1e-3, rtol=1e-4), ( + f"Torque mismatch with CoM offset + rotation.\n" + f"Expected:\n{expected_torque_b}\nGot:\n{composer.out_torque_b.numpy()}" + ) + + +# ============================================================================ +# Deprecated set_forces_and_torques Tests +# ============================================================================ + + +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +def test_deprecated_set_forces_and_torques_emits_warning(device: str): + """Test that the deprecated set_forces_and_torques wrapper emits a warning and works.""" + num_envs, num_bodies = 4, 2 + rng = np.random.default_rng(seed=80) + + mock_asset = create_mock_asset(num_envs, num_bodies, device) + composer = WrenchComposer(mock_asset) + + forces_np = rng.uniform(-50.0, 50.0, (num_envs, num_bodies, 3)).astype(np.float32) + + with pytest.warns(DeprecationWarning, match="set_forces_and_torques.*is deprecated"): + composer.set_forces_and_torques( + forces=wp.from_numpy(forces_np, dtype=wp.vec3f, device=device), + ) + + composer.compose_to_body_frame() + assert np.allclose(composer.out_force_b.numpy(), forces_np, atol=1e-4, rtol=1e-5) + + +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +def test_deprecated_set_forces_and_torques_clears_previous(device: str): + """Test that deprecated set_forces_and_torques actually replaces previous values.""" + num_envs, num_bodies = 4, 2 + rng = np.random.default_rng(seed=81) + + mock_asset = create_mock_asset(num_envs, num_bodies, device) + composer = WrenchComposer(mock_asset) + + # First add some forces + forces_a_np = rng.uniform(-50.0, 50.0, (num_envs, num_bodies, 3)).astype(np.float32) + composer.add_forces_and_torques_index( + forces=wp.from_numpy(forces_a_np, dtype=wp.vec3f, device=device), + ) + + # Then set via deprecated method -- should replace + forces_b_np = rng.uniform(-50.0, 50.0, (num_envs, num_bodies, 3)).astype(np.float32) + with pytest.warns(DeprecationWarning): + composer.set_forces_and_torques( + forces=wp.from_numpy(forces_b_np, dtype=wp.vec3f, device=device), + ) + + composer.compose_to_body_frame() + assert np.allclose(composer.out_force_b.numpy(), forces_b_np, atol=1e-4, rtol=1e-5), ( + "Deprecated set_forces_and_torques did not replace previous values" + ) diff --git a/source/isaaclab/test/utils/test_wrench_composer_integration.py b/source/isaaclab/test/utils/test_wrench_composer_integration.py new file mode 100644 index 000000000000..f71690d96541 --- /dev/null +++ b/source/isaaclab/test/utils/test_wrench_composer_integration.py @@ -0,0 +1,817 @@ +# 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 + +"""Integration tests for wrench composer with rigid objects. + +These tests validate that global forces/torques remain invariant under body rotation +""" + +"""Launch Isaac Sim Simulator first.""" + +from isaaclab.app import AppLauncher + +# launch omniverse app +simulation_app = AppLauncher(headless=True).app + +"""Rest everything follows.""" + +import pytest +import torch +import warp as wp + +import isaaclab.sim as sim_utils +from isaaclab.assets import RigidObject, RigidObjectCfg +from isaaclab.sim import build_simulation_context +from isaaclab.utils.assets import ISAAC_NUCLEUS_DIR + + +def generate_cubes_scene( + num_cubes: int = 1, + height: float = 1.0, + device: str = "cuda:0", +) -> tuple[RigidObject, torch.Tensor]: + """Generate a scene with the provided number of cubes.""" + origins = torch.tensor([(i * 1.0, 0, height) for i in range(num_cubes)]).to(device) + for i, origin in enumerate(origins): + sim_utils.create_prim(f"/World/Table_{i}", "Xform", translation=origin) + + spawn_cfg = sim_utils.UsdFileCfg( + usd_path=f"{ISAAC_NUCLEUS_DIR}/Props/Blocks/DexCube/dex_cube_instanceable.usd", + rigid_props=sim_utils.RigidBodyPropertiesCfg(), + ) + + cube_object_cfg = RigidObjectCfg( + prim_path="/World/Table_.*/Object", + spawn=spawn_cfg, + init_state=RigidObjectCfg.InitialStateCfg(pos=(0.0, 0.0, height)), + ) + cube_object = RigidObject(cfg=cube_object_cfg) + return cube_object, origins + + +N_STEPS = 100 +FORCE_MAGNITUDE = 10.0 +TORQUE_MAGNITUDE = 1.0 + + +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +def test_global_force_invariant_under_rotation(device): + """Test that a permanent global force produces the same acceleration before and after body rotation. + + A global +X force is applied. After 100 steps the body is rotated 180deg about Z. + The acceleration (delta_v per phase) should be the same in both phases because the + force is in the global frame and should not rotate with the body. + """ + with build_simulation_context(device=device, gravity_enabled=False, auto_add_lighting=True) as sim: + sim._app_control_on_stop_handle = None + cube_object, _ = generate_cubes_scene(num_cubes=1, device=device) + + sim.reset() + + body_ids, _ = cube_object.find_bodies(".*") + mass = float(wp.to_torch(cube_object.root_view.get_masses())[0]) + com = wp.to_torch(cube_object.data.body_com_pos_w).clone() + + # Apply permanent global force along +X at CoM + forces = torch.zeros(1, len(body_ids), 3, device=device) + forces[..., 0] = FORCE_MAGNITUDE + torques = torch.zeros(1, len(body_ids), 3, device=device) + + cube_object.permanent_wrench_composer.set_forces_and_torques( + forces=forces, + torques=torques, + positions=com, + body_ids=body_ids, + is_global=True, + ) + + # Phase 1: run N_STEPS + for _ in range(N_STEPS): + cube_object.write_data_to_sim() + sim.step() + cube_object.update(sim.cfg.dt) + + vel_after_phase1 = wp.to_torch(cube_object.data.root_lin_vel_w)[0].clone() + + # Rotate body 180deg about Z (quat wxyz = [0, 0, 0, 1]) while keeping velocity + root_pose = wp.to_torch(cube_object.data.root_state_w)[0, :7].clone().unsqueeze(0) + root_pose[0, 3:7] = torch.tensor([0.0, 0.0, 1.0, 0.0], device=device) # 180deg about Z (xyzw) + cube_object.write_root_pose_to_sim(root_pose) + + # Phase 2: run N_STEPS more + for _ in range(N_STEPS): + cube_object.write_data_to_sim() + sim.step() + cube_object.update(sim.cfg.dt) + + vel_after_phase2 = wp.to_torch(cube_object.data.root_lin_vel_w)[0].clone() + + # Acceleration should be same in both phases: delta_v_phase2 ≈ delta_v_phase1 + delta_v_phase1 = vel_after_phase1[0].item() # vx after phase 1 + delta_v_phase2 = vel_after_phase2[0].item() - vel_after_phase1[0].item() # vx gained in phase 2 + + expected_dv = FORCE_MAGNITUDE / mass * sim.cfg.dt * N_STEPS + + torch.testing.assert_close( + torch.tensor(delta_v_phase1), + torch.tensor(expected_dv), + rtol=0.001, + atol=0.0001, + ) + torch.testing.assert_close( + torch.tensor(delta_v_phase2), + torch.tensor(expected_dv), + rtol=0.001, + atol=0.0001, + ) + + # Y and Z velocity should remain ~0 + assert abs(vel_after_phase2[1].item()) < 0.5, f"Unexpected Y velocity: {vel_after_phase2[1].item()}" + assert abs(vel_after_phase2[2].item()) < 0.5, f"Unexpected Z velocity: {vel_after_phase2[2].item()}" + + +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +def test_local_force_follows_rotation(device): + """Test that a permanent local force rotates with the body. + + A local +X force is applied. After 100 steps the body is rotated 180deg about Z. + Since local +X is now world -X, the force should decelerate the body back towards zero velocity. + """ + with build_simulation_context(device=device, gravity_enabled=False, auto_add_lighting=True) as sim: + sim._app_control_on_stop_handle = None + cube_object, _ = generate_cubes_scene(num_cubes=1, device=device) + + sim.reset() + + body_ids, _ = cube_object.find_bodies(".*") + + # Apply permanent local force along body +X + forces = torch.zeros(1, len(body_ids), 3, device=device) + forces[..., 0] = FORCE_MAGNITUDE + torques = torch.zeros(1, len(body_ids), 3, device=device) + + cube_object.permanent_wrench_composer.set_forces_and_torques( + forces=forces, + torques=torques, + body_ids=body_ids, + is_global=False, + ) + + # Phase 1: run N_STEPS — object accelerates along world +X + for _ in range(N_STEPS): + cube_object.write_data_to_sim() + sim.step() + cube_object.update(sim.cfg.dt) + + vel_after_phase1 = wp.to_torch(cube_object.data.root_lin_vel_w)[0].clone() + assert vel_after_phase1[0].item() > 1.0, "Object should be moving in +X" + + # Rotate body 180deg about Z while keeping velocity + root_pose = wp.to_torch(cube_object.data.root_state_w)[0, :7].clone().unsqueeze(0) + root_pose[0, 3:7] = torch.tensor([0.0, 0.0, 1.0, 0.0], device=device) # 180deg about Z (xyzw) + cube_object.write_root_pose_to_sim(root_pose) + + # Phase 2: run N_STEPS — local +X is now world -X, so force decelerates + for _ in range(N_STEPS): + cube_object.write_data_to_sim() + sim.step() + cube_object.update(sim.cfg.dt) + + vel_after_phase2 = wp.to_torch(cube_object.data.root_lin_vel_w)[0].clone() + + # Velocity should be approximately zero: decelerated by the same amount as it accelerated + torch.testing.assert_close( + vel_after_phase2[0], + torch.tensor(0.0, device=device), + atol=0.0001, + rtol=0.001, + ) + + +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +def test_global_force_at_offset_generates_torque(device): + """Test that a global force applied at an offset from CoM generates the expected torque. + + A global +X force applied at +1m Y offset from CoM should produce: + - Linear acceleration in +X + - Angular acceleration about -Z (from cross product: (0,1,0) × (10,0,0) = (0,0,-10)) + """ + with build_simulation_context(device=device, gravity_enabled=False, auto_add_lighting=True) as sim: + sim._app_control_on_stop_handle = None + cube_object, _ = generate_cubes_scene(num_cubes=1, device=device) + + sim.reset() + + body_ids, _ = cube_object.find_bodies(".*") + + # Force at offset: +1m in Y from CoM (global frame) + forces = torch.zeros(1, len(body_ids), 3, device=device) + forces[..., 0] = FORCE_MAGNITUDE # +X force + + torques = torch.zeros(1, len(body_ids), 3, device=device) + + # Position offset: CoM position + 1m in Y (global frame) + com_pos = wp.to_torch(cube_object.data.body_com_pos_w)[:, body_ids, :3].clone() + positions = com_pos.clone() + positions[..., 1] += 1.0 # +1m Y offset + + cube_object.permanent_wrench_composer.set_forces_and_torques( + forces=forces, + torques=torques, + positions=positions, + body_ids=body_ids, + is_global=True, + ) + + # Run 50 steps + for _ in range(50): + cube_object.write_data_to_sim() + sim.step() + cube_object.update(sim.cfg.dt) + + lin_vel = wp.to_torch(cube_object.data.root_lin_vel_w)[0] + ang_vel = wp.to_torch(cube_object.data.root_ang_vel_w)[0] + + # Linear velocity in +X should be positive + assert lin_vel[0].item() > 0.1, f"Expected positive X velocity, got {lin_vel[0].item()}" + + # Angular velocity about Z should be negative (cross product: r × F, r=(0,1,0), F=(10,0,0) -> (0,0,-10)) + assert ang_vel[2].item() < -0.1, f"Expected negative Z angular velocity, got {ang_vel[2].item()}" + + +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +def test_global_torque_invariant_under_rotation(device): + """Test that a permanent global torque produces the same angular acceleration before and after rotation. + + A global +Z torque is applied. After 100 steps the body is rotated 90deg about X. + The angular acceleration (delta_omega per phase) about Z should be the same in both phases + because the torque is in the global frame. + """ + with build_simulation_context(device=device, gravity_enabled=False, auto_add_lighting=True) as sim: + sim._app_control_on_stop_handle = None + cube_object, _ = generate_cubes_scene(num_cubes=1, device=device) + + sim.reset() + + body_ids, _ = cube_object.find_bodies(".*") + + # Apply permanent global torque about +Z + forces = torch.zeros(1, len(body_ids), 3, device=device) + torques = torch.zeros(1, len(body_ids), 3, device=device) + torques[..., 2] = TORQUE_MAGNITUDE + + cube_object.permanent_wrench_composer.set_forces_and_torques( + forces=forces, + torques=torques, + body_ids=body_ids, + is_global=True, + ) + + # Phase 1: run N_STEPS + for _ in range(N_STEPS): + cube_object.write_data_to_sim() + sim.step() + cube_object.update(sim.cfg.dt) + + omega_z_after_phase1 = wp.to_torch(cube_object.data.root_ang_vel_w)[0, 2].clone().item() + + # Rotate body 90deg about X and zero out velocities so phase 2 starts from rest + # (avoids gyroscopic cross-coupling at high omega) + root_pose = wp.to_torch(cube_object.data.root_state_w)[0, :7].clone().unsqueeze(0) + root_pose[0, 3:7] = torch.tensor([0.7071, 0.0, 0.0, 0.7071], device=device) # 90deg about X (xyzw) + cube_object.write_root_pose_to_sim(root_pose) + cube_object.write_root_velocity_to_sim(torch.zeros(1, 6, device=device)) + + # Phase 2: run N_STEPS from rest with different body orientation + for _ in range(N_STEPS): + cube_object.write_data_to_sim() + sim.step() + cube_object.update(sim.cfg.dt) + + omega_z_after_phase2 = wp.to_torch(cube_object.data.root_ang_vel_w)[0, 2].clone().item() + + # Both phases start from rest — angular acceleration about Z should be the same + torch.testing.assert_close( + torch.tensor(omega_z_after_phase1), + torch.tensor(omega_z_after_phase2), + rtol=0.001, + atol=0.0001, + ) + + +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +def test_global_force_torque_after_translation(device): + """Test that global force torque updates dynamically when the body translates. + + Phase 1: Cube at (1,0,0). Global force F=(0,10,0) applied at explicit position (1,0,0). + stored_torque = cross((1,0,0), (0,10,0)) = (0,0,10) + correction = -cross((1,0,0), (0,10,0)) = (0,0,-10) + net torque = 0 → no rotation, only linear acceleration in +Y. + + Phase 2: Teleport cube to origin (0,0,0), zero velocity, don't re-apply force. + stored_torque = (0,0,10) (unchanged in buffer) + correction = -cross((0,0,0), (0,10,0)) = (0,0,0) + net torque = (0,0,10) → rotation about +Z. + """ + with build_simulation_context(device=device, gravity_enabled=False, auto_add_lighting=True) as sim: + sim._app_control_on_stop_handle = None + cube_object, _ = generate_cubes_scene(num_cubes=1, height=1.0, device=device) + + sim.reset() + + body_ids, _ = cube_object.find_bodies(".*") + + # Phase 1 setup: Move cube to (1, 0, 1) and apply force at (1, 0, 1) + root_state = wp.to_torch(cube_object.data.root_state_w).clone() + root_state[0, 0] = 1.0 # x = 1 + root_state[0, 1] = 0.0 # y = 0 + root_state[0, 2] = 1.0 # z = 1 + root_state[0, 3:7] = torch.tensor([0.0, 0.0, 0.0, 1.0], device=device) # identity quat (xyzw) + root_state[0, 7:] = 0.0 # zero velocity + cube_object.write_root_state_to_sim(root_state) + + # Step once to let the state settle + sim.step() + cube_object.update(sim.cfg.dt) + + # Get current CoM position for the force application point + com_pos = wp.to_torch(cube_object.data.body_com_pos_w)[:, body_ids, :3].clone() + + forces = torch.zeros(1, len(body_ids), 3, device=device) + forces[..., 1] = FORCE_MAGNITUDE # +Y force + torques = torch.zeros(1, len(body_ids), 3, device=device) + + cube_object.permanent_wrench_composer.set_forces_and_torques( + forces=forces, + torques=torques, + positions=com_pos, + body_ids=body_ids, + is_global=True, + ) + + # Phase 1: run 50 steps — force at CoM, expect no rotation + for _ in range(50): + cube_object.write_data_to_sim() + sim.step() + cube_object.update(sim.cfg.dt) + + ang_vel_phase1 = wp.to_torch(cube_object.data.root_ang_vel_w)[0].clone() + lin_vel_phase1 = wp.to_torch(cube_object.data.root_lin_vel_w)[0].clone() + + # Should have linear velocity in +Y + assert lin_vel_phase1[1].item() > 0.1, f"Expected positive Y velocity, got {lin_vel_phase1[1].item()}" + + # Angular velocity should be ~0 (force applied at CoM → no torque) + assert abs(ang_vel_phase1[2].item()) < 0.1, ( + f"Expected ~0 Z angular velocity in phase 1, got {ang_vel_phase1[2].item()}" + ) + + # Phase 2: Teleport cube to origin, zero velocity, don't re-apply force + root_state2 = wp.to_torch(cube_object.data.root_state_w).clone() + root_state2[0, 0] = 0.0 # x = 0 + root_state2[0, 1] = 0.0 + root_state2[0, 2] = 1.0 # z = 1 + root_state2[0, 3:7] = torch.tensor([0.0, 0.0, 0.0, 1.0], device=device) # identity (xyzw) + root_state2[0, 7:] = 0.0 # zero velocity + cube_object.write_root_state_to_sim(root_state2) + + # Step once to let state settle + sim.step() + cube_object.update(sim.cfg.dt) + + # Phase 2: run 50 steps — body at origin but stored torque = cross((1,0,1), (0,10,0)) = (-10,0,10) + # correction = -cross((0,0,1), (0,10,0)) = -(0,0,0 - but wait, z=1) + # Actually: stored = cross((com_x,com_y,com_z), (0,10,0)) + # After teleport: correction = -cross(new_pos, F), net torque ≠ 0 since positions differ + for _ in range(50): + cube_object.write_data_to_sim() + sim.step() + cube_object.update(sim.cfg.dt) + + ang_vel_phase2 = wp.to_torch(cube_object.data.root_ang_vel_w)[0].clone() + + # The X component of position changed from ~1 to ~0, so torque about Z changes. + # stored_torque_z = com_x * Fy = ~1 * 10 = ~10 + # After teleport, correction_z = -new_x * Fy = ~0 * 10 = ~0 + # net torque_z ≈ 10 → positive Z angular velocity + assert ang_vel_phase2[2].item() > 0.5, ( + f"Expected positive Z angular velocity in phase 2, got {ang_vel_phase2[2].item()}" + ) + + +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +def test_global_force_torque_reverses_on_opposite_side(device): + """Test that dynamic correction produces correct torque sign depending on body position. + + Phase 1: Cube at (-1, 0, 1). Global F=(0, 10, 0) at world point P=(0, 0, 1). + net torque_z = cross(P - link_pos, F)_z = cross((1,0,0), (0,10,0))_z = +10 + → positive Z angular velocity + + Phase 2: Teleport cube to (+1, 0, 1), zero velocity, don't re-apply force. + net torque_z = cross(P - link_pos, F)_z = cross((-1,0,0), (0,10,0))_z = -10 + → negative Z angular velocity + """ + with build_simulation_context(device=device, gravity_enabled=False, auto_add_lighting=True) as sim: + sim._app_control_on_stop_handle = None + cube_object, _ = generate_cubes_scene(num_cubes=1, height=1.0, device=device) + + sim.reset() + + body_ids, _ = cube_object.find_bodies(".*") + + # Move cube to (-1, 0, 1) + root_state = wp.to_torch(cube_object.data.root_state_w).clone() + root_state[0, 0] = -1.0 + root_state[0, 1] = 0.0 + root_state[0, 2] = 1.0 + root_state[0, 3:7] = torch.tensor([0.0, 0.0, 0.0, 1.0], device=device) # identity (xyzw) + root_state[0, 7:] = 0.0 + cube_object.write_root_state_to_sim(root_state) + sim.step() + cube_object.update(sim.cfg.dt) + + # Apply permanent global F=(0, 10, 0) at world point P=(0, 0, 1) + forces = torch.zeros(1, len(body_ids), 3, device=device) + forces[..., 1] = FORCE_MAGNITUDE + torques = torch.zeros(1, len(body_ids), 3, device=device) + positions = torch.zeros(1, len(body_ids), 3, device=device) + positions[..., 2] = 1.0 # P = (0, 0, 1) + + cube_object.permanent_wrench_composer.set_forces_and_torques( + forces=forces, + torques=torques, + positions=positions, + body_ids=body_ids, + is_global=True, + ) + + # Phase 1: run 50 steps — expect positive Z angular velocity + for _ in range(50): + cube_object.write_data_to_sim() + sim.step() + cube_object.update(sim.cfg.dt) + + omega_z_phase1 = wp.to_torch(cube_object.data.root_ang_vel_w)[0, 2].item() + assert omega_z_phase1 > 0.1, f"Phase 1: expected positive omega_z, got {omega_z_phase1}" + + # Phase 2: Teleport cube to (+1, 0, 1), zero velocity + root_state2 = wp.to_torch(cube_object.data.root_state_w).clone() + root_state2[0, 0] = 1.0 + root_state2[0, 1] = 0.0 + root_state2[0, 2] = 1.0 + root_state2[0, 3:7] = torch.tensor([0.0, 0.0, 0.0, 1.0], device=device) # identity (xyzw) + root_state2[0, 7:] = 0.0 + cube_object.write_root_state_to_sim(root_state2) + sim.step() + cube_object.update(sim.cfg.dt) + + # Phase 2: run 50 steps — expect negative Z angular velocity + for _ in range(50): + cube_object.write_data_to_sim() + sim.step() + cube_object.update(sim.cfg.dt) + + omega_z_phase2 = wp.to_torch(cube_object.data.root_ang_vel_w)[0, 2].item() + assert omega_z_phase2 < -0.1, f"Phase 2: expected negative omega_z, got {omega_z_phase2}" + + +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +def test_global_force_no_position_no_torque(device): + """Test that global force without positions produces no torque (applied at CoM). + + A body at (2, 0, 1) with global F=(0, 10, 0) and no positions should experience + only linear acceleration, no rotation. The force is applied at the body's CoM. + """ + with build_simulation_context(device=device, gravity_enabled=False, auto_add_lighting=True) as sim: + sim._app_control_on_stop_handle = None + cube_object, _ = generate_cubes_scene(num_cubes=1, height=1.0, device=device) + + sim.reset() + + body_ids, _ = cube_object.find_bodies(".*") + + # Move cube to (2, 0, 1) + root_state = wp.to_torch(cube_object.data.root_state_w).clone() + root_state[0, 0] = 2.0 + root_state[0, 1] = 0.0 + root_state[0, 2] = 1.0 + root_state[0, 3:7] = torch.tensor([0.0, 0.0, 0.0, 1.0], device=device) # identity (xyzw) + root_state[0, 7:] = 0.0 + cube_object.write_root_state_to_sim(root_state) + sim.step() + cube_object.update(sim.cfg.dt) + + # Apply global F=(0, 10, 0) WITHOUT positions → force at CoM, no torque + forces = torch.zeros(1, len(body_ids), 3, device=device) + forces[..., 1] = FORCE_MAGNITUDE + torques = torch.zeros(1, len(body_ids), 3, device=device) + + cube_object.permanent_wrench_composer.set_forces_and_torques( + forces=forces, + torques=torques, + body_ids=body_ids, + is_global=True, + ) + + # Run 50 steps + for _ in range(50): + cube_object.write_data_to_sim() + sim.step() + cube_object.update(sim.cfg.dt) + + omega_z = wp.to_torch(cube_object.data.root_ang_vel_w)[0, 2].item() + # No positions → force at CoM → zero torque → zero angular velocity + assert abs(omega_z) < 0.01, f"Expected ~zero omega_z for force at CoM, got {omega_z}" + + # Should still have linear acceleration in +Y + lin_vel_y = wp.to_torch(cube_object.data.root_lin_vel_w)[0, 1].item() + assert lin_vel_y > 0.1, f"Expected positive Y velocity from applied force, got {lin_vel_y}" + + +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +def test_multi_cube_different_torques_from_same_force(device): + """Test kernel indexing across multiple envs with different CoM positions. + + 2 cubes: Cube 0 at (-1, 0, 1), Cube 1 at (+1, 0, 1). + Same global F=(0, 10, 0) at same world point P=(0, 0, 1) to both cubes. + Cube 0: torque_z = cross((1,0,0), (0,10,0))_z = +10 → omega_z > 0 + Cube 1: torque_z = cross((-1,0,0), (0,10,0))_z = -10 → omega_z < 0 + Both have same linear acceleration in +Y. + """ + with build_simulation_context(device=device, gravity_enabled=False, auto_add_lighting=True) as sim: + sim._app_control_on_stop_handle = None + cube_object, _ = generate_cubes_scene(num_cubes=2, height=1.0, device=device) + + sim.reset() + + body_ids, _ = cube_object.find_bodies(".*") + + # Position cubes: Cube 0 at (-1, 0, 1), Cube 1 at (+1, 0, 1) + root_state = wp.to_torch(cube_object.data.root_state_w).clone() + root_state[0, 0] = -1.0 + root_state[0, 1] = 0.0 + root_state[0, 2] = 1.0 + root_state[0, 3:7] = torch.tensor([0.0, 0.0, 0.0, 1.0], device=device) # identity (xyzw) + root_state[0, 7:] = 0.0 + + root_state[1, 0] = 1.0 + root_state[1, 1] = 0.0 + root_state[1, 2] = 1.0 + root_state[1, 3:7] = torch.tensor([0.0, 0.0, 0.0, 1.0], device=device) # identity (xyzw) + root_state[1, 7:] = 0.0 + cube_object.write_root_state_to_sim(root_state) + sim.step() + cube_object.update(sim.cfg.dt) + + # Apply same global F=(0, 10, 0) at P=(0, 0, 1) to both cubes + forces = torch.zeros(2, len(body_ids), 3, device=device) + forces[..., 1] = FORCE_MAGNITUDE + torques = torch.zeros(2, len(body_ids), 3, device=device) + positions = torch.zeros(2, len(body_ids), 3, device=device) + positions[..., 2] = 1.0 # P = (0, 0, 1) + + cube_object.permanent_wrench_composer.set_forces_and_torques( + forces=forces, + torques=torques, + positions=positions, + body_ids=body_ids, + is_global=True, + ) + + # Run 50 steps + for _ in range(50): + cube_object.write_data_to_sim() + sim.step() + cube_object.update(sim.cfg.dt) + + # Cube 0: omega_z > 0 (force point is to the right of CoM) + omega_z_0 = wp.to_torch(cube_object.data.root_ang_vel_w)[0, 2].item() + assert omega_z_0 > 0.1, f"Cube 0: expected positive omega_z, got {omega_z_0}" + + # Cube 1: omega_z < 0 (force point is to the left of CoM) + omega_z_1 = wp.to_torch(cube_object.data.root_ang_vel_w)[1, 2].item() + assert omega_z_1 < -0.1, f"Cube 1: expected negative omega_z, got {omega_z_1}" + + # Both cubes should have same linear velocity in +Y (same force magnitude) + lin_vel_y_0 = wp.to_torch(cube_object.data.root_lin_vel_w)[0, 1].item() + lin_vel_y_1 = wp.to_torch(cube_object.data.root_lin_vel_w)[1, 1].item() + assert abs(lin_vel_y_0 - lin_vel_y_1) < 0.5, ( + f"Both cubes should have similar Y velocity, got {lin_vel_y_0} and {lin_vel_y_1}" + ) + + +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +def test_global_force_torque_far_from_origin(device): + """Test that global force torque correction produces correct physics at large world coordinates. + + Two cubes with identical relative geometry (force offset = (1, 0, 0) from CoM): + Cube 0 at (0, 0, 1) — near origin (reference) + Cube 1 at (2000, 0, 1) — far from origin + + Both get global F=(0, 10, 0) at offset (1, 0, 0) from their respective CoMs. + Expected torque: cross((1,0,0), (0,10,0)) = (0, 0, 10) for both. + + The compose kernel computes cross(P, F) - cross(link_pos, F): + Cube 0: cross((1,0,1), F) - cross((0,0,1), F) — small values, no cancellation + Cube 1: cross((2001,0,1), F) - cross((2000,0,1), F) — large values nearly cancel + + Both cubes should produce the same angular and linear velocities. + """ + with build_simulation_context(device=device, gravity_enabled=False, auto_add_lighting=True) as sim: + sim._app_control_on_stop_handle = None + cube_object, _ = generate_cubes_scene(num_cubes=2, height=1.0, device=device) + + sim.reset() + + body_ids, _ = cube_object.find_bodies(".*") + + # Position cubes: Cube 0 near origin, Cube 1 far from origin + root_state = wp.to_torch(cube_object.data.root_state_w).clone() + # Cube 0 at (0, 0, 1) + root_state[0, 0] = 0.0 + root_state[0, 1] = 0.0 + root_state[0, 2] = 1.0 + root_state[0, 3:7] = torch.tensor([0.0, 0.0, 0.0, 1.0], device=device) # identity (xyzw) + root_state[0, 7:] = 0.0 + # Cube 1 at (2000, 0, 1) + root_state[1, 0] = 2000.0 + root_state[1, 1] = 0.0 + root_state[1, 2] = 1.0 + root_state[1, 3:7] = torch.tensor([0.0, 0.0, 0.0, 1.0], device=device) # identity (xyzw) + root_state[1, 7:] = 0.0 + cube_object.write_root_state_to_sim(root_state) + sim.step() + cube_object.update(sim.cfg.dt) + + # Apply F=(0, 10, 0) at +1m X offset from each cube's CoM + forces = torch.zeros(2, len(body_ids), 3, device=device) + forces[..., 1] = FORCE_MAGNITUDE # +Y force + torques = torch.zeros(2, len(body_ids), 3, device=device) + + # Positions: each cube's CoM + (1, 0, 0) + com_pos = wp.to_torch(cube_object.data.body_com_pos_w)[:, body_ids, :3].clone() + positions = com_pos.clone() + positions[..., 0] += 1.0 # +1m X offset from CoM + + cube_object.permanent_wrench_composer.set_forces_and_torques( + forces=forces, + torques=torques, + positions=positions, + body_ids=body_ids, + is_global=True, + ) + + # Run 50 steps + for _ in range(50): + cube_object.write_data_to_sim() + sim.step() + cube_object.update(sim.cfg.dt) + + # Both cubes should have positive omega_z (cross((1,0,0), (0,10,0)) = (0,0,10)) + omega_z_0 = wp.to_torch(cube_object.data.root_ang_vel_w)[0, 2].item() + omega_z_1 = wp.to_torch(cube_object.data.root_ang_vel_w)[1, 2].item() + assert omega_z_0 > 0.1, f"Cube 0: expected positive omega_z, got {omega_z_0}" + assert omega_z_1 > 0.1, f"Cube 1: expected positive omega_z, got {omega_z_1}" + + # omega_z values should match within 1% (same relative geometry) + torch.testing.assert_close( + torch.tensor(omega_z_0), + torch.tensor(omega_z_1), + rtol=0.01, + atol=0.0, + msg=lambda msg: ( + f"Angular velocity mismatch between near-origin and far-from-origin cubes:\n" + f" Cube 0 (near): omega_z = {omega_z_0:.6f}\n" + f" Cube 1 (far): omega_z = {omega_z_1:.6f}\n{msg}" + ), + ) + + # Linear velocity in +Y should also match + lin_vel_y_0 = wp.to_torch(cube_object.data.root_lin_vel_w)[0, 1].item() + lin_vel_y_1 = wp.to_torch(cube_object.data.root_lin_vel_w)[1, 1].item() + torch.testing.assert_close( + torch.tensor(lin_vel_y_0), + torch.tensor(lin_vel_y_1), + rtol=0.01, + atol=0.0, + msg=lambda msg: ( + f"Linear velocity mismatch between near-origin and far-from-origin cubes:\n" + f" Cube 0 (near): lin_vel_y = {lin_vel_y_0:.6f}\n" + f" Cube 1 (far): lin_vel_y = {lin_vel_y_1:.6f}\n{msg}" + ), + ) + + +@pytest.mark.parametrize("device", ["cuda:0"]) +def test_global_force_no_position_no_rotation_large_offset(device): + """Test that a global force without positions produces no rotation at large offsets. + + A cube is placed at (2000, 0, 1) and a global force F=(0, 10, 0) is applied + without positions. The cube should accelerate linearly but not rotate. + Before the fix, this would produce torque proportional to 2000 and cause rotation. + """ + with build_simulation_context( + device=device, add_ground_plane=False, auto_add_lighting=True, gravity_enabled=False + ) as sim: + sim._app_control_on_stop_handle = None + cube_object, _ = generate_cubes_scene(num_cubes=1, height=1.0, device=device) + + sim.reset() + + body_ids, _ = cube_object.find_bodies(".*") + + # Place cube at large X offset + root_state = wp.to_torch(cube_object.data.default_root_state).clone() + root_state[0, 0] = 2000.0 # large X position + root_state[0, 1] = 0.0 + root_state[0, 2] = 1.0 + cube_object.write_root_pose_to_sim(root_state[:, :7]) + cube_object.write_root_velocity_to_sim(root_state[:, 7:]) + cube_object.reset() + + # Apply global force without positions (should go to CoM, no torque) + forces = torch.zeros(cube_object.num_instances, len(body_ids), 3, device=device) + forces[0, :, 1] = 10.0 # F_y = 10 N + + cube_object.permanent_wrench_composer.set_forces_and_torques( + forces=forces, + body_ids=body_ids, + is_global=True, + ) + + # Step simulation + for _ in range(50): + cube_object.write_data_to_sim() + sim.step() + cube_object.update(sim.cfg.dt) + + # Check: angular velocity should be near zero (no rotation) + ang_vel = wp.to_torch(cube_object.data.root_ang_vel_w)[0] + assert torch.allclose(ang_vel, torch.zeros(3, device=device), atol=0.01), ( + f"Expected near-zero angular velocity, got {ang_vel}. " + "Global force without positions should not produce torque." + ) + + # Check: linear velocity in Y should be positive (force is in +Y) + lin_vel = wp.to_torch(cube_object.data.root_lin_vel_w)[0] + assert lin_vel[1] > 0.1, f"Expected positive Y velocity from applied force, got {lin_vel[1]}" + + +@pytest.mark.parametrize("device", ["cuda:0"]) +def test_global_force_at_com_position_no_rotation_large_offset(device): + """Test that a global force with position at CoM produces no rotation at large offsets. + + A cube is placed at (2000, 0, 1) and a global force F=(0, 10, 0) is applied + at the cube's position (i.e., at its CoM). This should produce zero torque, + serving as a control test alongside test_global_force_no_position_no_rotation_large_offset. + """ + with build_simulation_context( + device=device, add_ground_plane=False, auto_add_lighting=True, gravity_enabled=False + ) as sim: + sim._app_control_on_stop_handle = None + cube_object, _ = generate_cubes_scene(num_cubes=1, height=1.0, device=device) + + sim.reset() + + body_ids, _ = cube_object.find_bodies(".*") + + # Place cube at large X offset + root_state = wp.to_torch(cube_object.data.default_root_state).clone() + root_state[0, 0] = 2000.0 + root_state[0, 1] = 0.0 + root_state[0, 2] = 1.0 + cube_object.write_root_pose_to_sim(root_state[:, :7]) + cube_object.write_root_velocity_to_sim(root_state[:, 7:]) + cube_object.reset() + + # Apply global force AT the cube's position (torque should cancel) + forces = torch.zeros(cube_object.num_instances, len(body_ids), 3, device=device) + forces[0, :, 1] = 10.0 + + positions = torch.zeros(cube_object.num_instances, len(body_ids), 3, device=device) + positions[0, :, 0] = 2000.0 + positions[0, :, 2] = 1.0 + + cube_object.permanent_wrench_composer.set_forces_and_torques( + forces=forces, + positions=positions, + body_ids=body_ids, + is_global=True, + ) + + for _ in range(50): + cube_object.write_data_to_sim() + sim.step() + cube_object.update(sim.cfg.dt) + + # Force at CoM → no rotation + ang_vel = wp.to_torch(cube_object.data.root_ang_vel_w)[0] + assert torch.allclose(ang_vel, torch.zeros(3, device=device), atol=0.01), ( + f"Expected near-zero angular velocity, got {ang_vel}. " + "Global force at CoM position should not produce torque." + ) + + lin_vel = wp.to_torch(cube_object.data.root_lin_vel_w)[0] + assert lin_vel[1] > 0.1, f"Expected positive Y velocity from applied force, got {lin_vel[1]}" diff --git a/source/isaaclab/test/utils/test_wrench_composer_vs_physx.py b/source/isaaclab/test/utils/test_wrench_composer_vs_physx.py new file mode 100644 index 000000000000..bb0a69132890 --- /dev/null +++ b/source/isaaclab/test/utils/test_wrench_composer_vs_physx.py @@ -0,0 +1,837 @@ +# 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 + +"""Integration tests comparing WrenchComposer output vs raw PhysX apply_forces_and_torques_at_position. + +Two identical rigid objects are placed in the same scene. One uses the WrenchComposer path +(set_forces_and_torques → write_data_to_sim → compose → PhysX apply with is_global=False), +the other uses the raw PhysX API directly (apply_forces_and_torques_at_position with matching +is_global flag). After N steps, both objects should have identical velocities. +""" + +"""Launch Isaac Sim Simulator first.""" + +from isaaclab.app import AppLauncher + +# launch omniverse app +simulation_app = AppLauncher(headless=True).app + +"""Rest everything follows.""" + +import math + +import pytest +import torch +import warp as wp + +import isaaclab.sim as sim_utils +from isaaclab.assets import RigidObject, RigidObjectCfg +from isaaclab.sim import build_simulation_context +from isaaclab.utils.assets import ISAAC_NUCLEUS_DIR + + +def generate_dual_cube_scene( + num_cubes: int = 1, + height: float = 1.0, + device: str = "cuda:0", + initial_rot: tuple[float, ...] | None = None, + spacing: float = 2.0, +) -> tuple[RigidObject, RigidObject]: + """Generate a scene with two sets of cubes: one for the composer path, one for raw PhysX. + + Both sets share the same spawn config and initial state (except a Y offset to avoid overlap). + + Args: + num_cubes: Number of cubes per group (environments). + height: Spawn height. + device: Simulation device. + initial_rot: Initial quaternion (x, y, z, w). Defaults to identity. + spacing: Distance between env origins in X. Defaults to 2.0. + + Returns: + Tuple of (cube_composer, cube_raw) RigidObject instances. + """ + if initial_rot is None: + initial_rot = (0.0, 0.0, 0.0, 1.0) # identity in (x,y,z,w) + + y_offset = max(spacing, 3.0) + + # Create Xform prims for both groups + for i in range(num_cubes): + origin_composer = (i * spacing, 0.0, height) + origin_raw = (i * spacing, y_offset, height) # Y offset to avoid overlap + sim_utils.create_prim(f"/World/Composer_{i}", "Xform", translation=origin_composer) + sim_utils.create_prim(f"/World/Raw_{i}", "Xform", translation=origin_raw) + + spawn_cfg = sim_utils.UsdFileCfg( + usd_path=f"{ISAAC_NUCLEUS_DIR}/Props/Blocks/DexCube/dex_cube_instanceable.usd", + rigid_props=sim_utils.RigidBodyPropertiesCfg(), + ) + + cube_composer_cfg = RigidObjectCfg( + prim_path="/World/Composer_.*/Object", + spawn=spawn_cfg, + init_state=RigidObjectCfg.InitialStateCfg(pos=(0.0, 0.0, height), rot=initial_rot), + ) + cube_composer = RigidObject(cfg=cube_composer_cfg) + + cube_raw_cfg = RigidObjectCfg( + prim_path="/World/Raw_.*/Object", + spawn=spawn_cfg, + init_state=RigidObjectCfg.InitialStateCfg(pos=(0.0, y_offset, height), rot=initial_rot), + ) + cube_raw = RigidObject(cfg=cube_raw_cfg) + + return cube_composer, cube_raw + + +N_STEPS = 50 +FORCE_MAGNITUDE = 10.0 +TORQUE_MAGNITUDE = 1.0 +# 45 degrees about Z: (cos(22.5°), 0, 0, sin(22.5°)) +ROT_45_Z = (0.0, 0.0, math.sin(math.pi / 8), math.cos(math.pi / 8)) # 45deg about Z in (x,y,z,w) + + +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +def test_composer_vs_physx_local_force(device): + """Baseline: local force at identity orientation. Composer and raw PhysX should match exactly.""" + with build_simulation_context(device=device, gravity_enabled=False, auto_add_lighting=True) as sim: + sim._app_control_on_stop_handle = None + cube_composer, cube_raw = generate_dual_cube_scene(num_cubes=1, device=device) + + sim.reset() + + body_ids, _ = cube_composer.find_bodies(".*") + + # Composer path: local force +X + forces = torch.zeros(1, len(body_ids), 3, device=device) + forces[..., 0] = FORCE_MAGNITUDE + torques = torch.zeros(1, len(body_ids), 3, device=device) + + cube_composer.permanent_wrench_composer.set_forces_and_torques( + forces=forces, + torques=torques, + body_ids=body_ids, + is_global=False, + ) + + # Raw PhysX data (flattened for PhysX view API) + raw_forces = torch.zeros(1, 3, device=device) + raw_forces[:, 0] = FORCE_MAGNITUDE + raw_torques = torch.zeros(1, 3, device=device) + raw_indices = cube_raw._ALL_INDICES + + for _ in range(N_STEPS): + cube_composer.write_data_to_sim() + cube_raw.write_data_to_sim() # no-op (composer inactive) + cube_raw.root_view.apply_forces_and_torques_at_position( + force_data=wp.from_torch(raw_forces.contiguous(), dtype=wp.float32), + torque_data=wp.from_torch(raw_torques.contiguous(), dtype=wp.float32), + position_data=None, + indices=raw_indices, + is_global=False, + ) + sim.step() + cube_composer.update(sim.cfg.dt) + cube_raw.update(sim.cfg.dt) + + # Compare velocities + torch.testing.assert_close( + wp.to_torch(cube_composer.data.root_lin_vel_w), + wp.to_torch(cube_raw.data.root_lin_vel_w), + rtol=1e-4, + atol=1e-4, + ) + # Both should have ~zero angular velocity (force at CoM, no torque) + torch.testing.assert_close( + wp.to_torch(cube_composer.data.root_ang_vel_w), + torch.zeros(1, 3, device=device), + rtol=0.0, + atol=1e-4, + ) + torch.testing.assert_close( + wp.to_torch(cube_raw.data.root_ang_vel_w), + torch.zeros(1, 3, device=device), + rtol=0.0, + atol=1e-4, + ) + + +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +def test_composer_vs_physx_global_force(device): + """Global force with non-identity rotation (45 deg Z). Rotation matters for frame conversion.""" + with build_simulation_context(device=device, gravity_enabled=False, auto_add_lighting=True) as sim: + sim._app_control_on_stop_handle = None + cube_composer, cube_raw = generate_dual_cube_scene(num_cubes=1, device=device, initial_rot=ROT_45_Z) + + sim.reset() + + body_ids, _ = cube_composer.find_bodies(".*") + + # Composer path: global force +X + forces = torch.zeros(1, len(body_ids), 3, device=device) + forces[..., 0] = FORCE_MAGNITUDE + torques = torch.zeros(1, len(body_ids), 3, device=device) + + cube_composer.permanent_wrench_composer.set_forces_and_torques( + forces=forces, + torques=torques, + body_ids=body_ids, + is_global=True, + ) + + # Raw PhysX data + raw_forces = torch.zeros(1, 3, device=device) + raw_forces[:, 0] = FORCE_MAGNITUDE + raw_torques = torch.zeros(1, 3, device=device) + raw_indices = cube_raw._ALL_INDICES + + for _ in range(N_STEPS): + cube_composer.write_data_to_sim() + cube_raw.write_data_to_sim() + cube_raw.root_view.apply_forces_and_torques_at_position( + force_data=wp.from_torch(raw_forces.contiguous(), dtype=wp.float32), + torque_data=wp.from_torch(raw_torques.contiguous(), dtype=wp.float32), + position_data=None, + indices=raw_indices, + is_global=True, + ) + sim.step() + cube_composer.update(sim.cfg.dt) + cube_raw.update(sim.cfg.dt) + + # Linear velocities should match (same global force, same mass) + torch.testing.assert_close( + wp.to_torch(cube_composer.data.root_lin_vel_w), + wp.to_torch(cube_raw.data.root_lin_vel_w), + rtol=1e-4, + atol=1e-4, + ) + # Angular velocities should match + torch.testing.assert_close( + wp.to_torch(cube_composer.data.root_ang_vel_w), + wp.to_torch(cube_raw.data.root_ang_vel_w), + rtol=1e-4, + atol=1e-4, + ) + # Both should have ~zero angular velocity (force at CoM, no torque) + torch.testing.assert_close( + wp.to_torch(cube_composer.data.root_ang_vel_w), + torch.zeros(1, 3, device=device), + rtol=0.0, + atol=1e-4, + ) + torch.testing.assert_close( + wp.to_torch(cube_raw.data.root_ang_vel_w), + torch.zeros(1, 3, device=device), + rtol=0.0, + atol=1e-4, + ) + + +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +def test_composer_vs_physx_local_force_at_position(device): + """Local force at a local offset. Both paths should produce identical cross-product torque.""" + with build_simulation_context(device=device, gravity_enabled=False, auto_add_lighting=True) as sim: + sim._app_control_on_stop_handle = None + cube_composer, cube_raw = generate_dual_cube_scene(num_cubes=1, device=device) + + sim.reset() + + body_ids, _ = cube_composer.find_bodies(".*") + + # Local force +X at local offset +0.5m Y + forces = torch.zeros(1, len(body_ids), 3, device=device) + forces[..., 0] = FORCE_MAGNITUDE + torques = torch.zeros(1, len(body_ids), 3, device=device) + positions = torch.zeros(1, len(body_ids), 3, device=device) + positions[..., 1] = 0.5 # +0.5m Y offset in local frame + + cube_composer.permanent_wrench_composer.set_forces_and_torques( + forces=forces, + torques=torques, + positions=positions, + body_ids=body_ids, + is_global=False, + ) + + # Raw PhysX data (local force at local position) + raw_forces = torch.zeros(1, 3, device=device) + raw_forces[:, 0] = FORCE_MAGNITUDE + raw_torques = torch.zeros(1, 3, device=device) + raw_positions = torch.zeros(1, 3, device=device) + raw_positions[:, 1] = 0.5 + raw_indices = cube_raw._ALL_INDICES + + for _ in range(N_STEPS): + cube_composer.write_data_to_sim() + cube_raw.write_data_to_sim() + cube_raw.root_view.apply_forces_and_torques_at_position( + force_data=wp.from_torch(raw_forces.contiguous(), dtype=wp.float32), + torque_data=wp.from_torch(raw_torques.contiguous(), dtype=wp.float32), + position_data=wp.from_torch(raw_positions.contiguous(), dtype=wp.float32), + indices=raw_indices, + is_global=False, + ) + sim.step() + cube_composer.update(sim.cfg.dt) + cube_raw.update(sim.cfg.dt) + + # Both linear and angular velocities should match + torch.testing.assert_close( + wp.to_torch(cube_composer.data.root_lin_vel_w), + wp.to_torch(cube_raw.data.root_lin_vel_w), + rtol=1e-4, + atol=1e-4, + ) + torch.testing.assert_close( + wp.to_torch(cube_composer.data.root_ang_vel_w), + wp.to_torch(cube_raw.data.root_ang_vel_w), + rtol=1e-4, + atol=1e-4, + ) + + # Sanity: angular velocity should be nonzero (cross-product torque) + assert torch.abs(wp.to_torch(cube_composer.data.root_ang_vel_w)[0, 2]).item() > 0.1, ( + "Expected nonzero Z angular velocity from cross-product torque" + ) + + +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +def test_composer_vs_physx_global_force_at_position(device): + """Global force at world position with non-identity rotation. Both rotation AND position correction matter.""" + with build_simulation_context(device=device, gravity_enabled=False, auto_add_lighting=True) as sim: + sim._app_control_on_stop_handle = None + cube_composer, cube_raw = generate_dual_cube_scene(num_cubes=1, device=device, initial_rot=ROT_45_Z) + + sim.reset() + + body_ids, _ = cube_composer.find_bodies(".*") + + # Global force +X + forces = torch.zeros(1, len(body_ids), 3, device=device) + forces[..., 0] = FORCE_MAGNITUDE + torques = torch.zeros(1, len(body_ids), 3, device=device) + + # Position = each cube's link_pos + offset (same offset for both) + offset = torch.zeros(1, len(body_ids), 3, device=device) + offset[..., 1] = 1.0 # +1m Y offset in world frame + + pos_composer = wp.to_torch(cube_composer.data.body_com_pos_w)[:, body_ids, :3].clone() + offset + pos_raw = wp.to_torch(cube_raw.data.body_com_pos_w)[:, body_ids, :3].clone() + offset + + cube_composer.permanent_wrench_composer.set_forces_and_torques( + forces=forces, + torques=torques, + positions=pos_composer, + body_ids=body_ids, + is_global=True, + ) + + # Raw PhysX data + raw_forces = torch.zeros(1, 3, device=device) + raw_forces[:, 0] = FORCE_MAGNITUDE + raw_torques = torch.zeros(1, 3, device=device) + raw_positions = pos_raw.view(-1, 3) + raw_indices = cube_raw._ALL_INDICES + + for _ in range(N_STEPS): + cube_composer.write_data_to_sim() + cube_raw.write_data_to_sim() + cube_raw.root_view.apply_forces_and_torques_at_position( + force_data=wp.from_torch(raw_forces.contiguous(), dtype=wp.float32), + torque_data=wp.from_torch(raw_torques.contiguous(), dtype=wp.float32), + position_data=wp.from_torch(raw_positions.contiguous(), dtype=wp.float32), + indices=raw_indices, + is_global=True, + ) + sim.step() + cube_composer.update(sim.cfg.dt) + cube_raw.update(sim.cfg.dt) + + # Both linear and angular velocities should match + torch.testing.assert_close( + wp.to_torch(cube_composer.data.root_lin_vel_w), + wp.to_torch(cube_raw.data.root_lin_vel_w), + rtol=1e-4, + atol=1e-4, + ) + torch.testing.assert_close( + wp.to_torch(cube_composer.data.root_ang_vel_w), + wp.to_torch(cube_raw.data.root_ang_vel_w), + rtol=1e-4, + atol=1e-4, + ) + + # Sanity: angular velocity should be nonzero (cross-product torque) + assert torch.abs(wp.to_torch(cube_composer.data.root_ang_vel_w)[0, 2]).item() > 0.1, ( + "Expected nonzero Z angular velocity from positional torque" + ) + + +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +def test_composer_vs_physx_local_torque(device): + """Local torque at identity orientation. Should produce matching angular velocity.""" + with build_simulation_context(device=device, gravity_enabled=False, auto_add_lighting=True) as sim: + sim._app_control_on_stop_handle = None + cube_composer, cube_raw = generate_dual_cube_scene(num_cubes=1, device=device) + + sim.reset() + + body_ids, _ = cube_composer.find_bodies(".*") + + # Composer path: local torque about +Z + forces = torch.zeros(1, len(body_ids), 3, device=device) + torques = torch.zeros(1, len(body_ids), 3, device=device) + torques[..., 2] = TORQUE_MAGNITUDE + + cube_composer.permanent_wrench_composer.set_forces_and_torques( + forces=forces, + torques=torques, + body_ids=body_ids, + is_global=False, + ) + + # Raw PhysX data + raw_forces = torch.zeros(1, 3, device=device) + raw_torques = torch.zeros(1, 3, device=device) + raw_torques[:, 2] = TORQUE_MAGNITUDE + raw_indices = cube_raw._ALL_INDICES + + for _ in range(N_STEPS): + cube_composer.write_data_to_sim() + cube_raw.write_data_to_sim() + cube_raw.root_view.apply_forces_and_torques_at_position( + force_data=wp.from_torch(raw_forces.contiguous(), dtype=wp.float32), + torque_data=wp.from_torch(raw_torques.contiguous(), dtype=wp.float32), + position_data=None, + indices=raw_indices, + is_global=False, + ) + sim.step() + cube_composer.update(sim.cfg.dt) + cube_raw.update(sim.cfg.dt) + + # Angular velocities should match + torch.testing.assert_close( + wp.to_torch(cube_composer.data.root_ang_vel_w), + wp.to_torch(cube_raw.data.root_ang_vel_w), + rtol=1e-4, + atol=1e-4, + ) + # Linear velocity should be ~zero for both (no force) + torch.testing.assert_close( + wp.to_torch(cube_composer.data.root_lin_vel_w), + torch.zeros(1, 3, device=device), + rtol=0.0, + atol=1e-4, + ) + torch.testing.assert_close( + wp.to_torch(cube_raw.data.root_lin_vel_w), + torch.zeros(1, 3, device=device), + rtol=0.0, + atol=1e-4, + ) + + +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +def test_composer_vs_physx_global_torque(device): + """Global torque with non-identity rotation (45 deg Z). Composer rotates to body frame internally.""" + with build_simulation_context(device=device, gravity_enabled=False, auto_add_lighting=True) as sim: + sim._app_control_on_stop_handle = None + cube_composer, cube_raw = generate_dual_cube_scene(num_cubes=1, device=device, initial_rot=ROT_45_Z) + + sim.reset() + + body_ids, _ = cube_composer.find_bodies(".*") + + # Composer path: global torque about +Z + forces = torch.zeros(1, len(body_ids), 3, device=device) + torques = torch.zeros(1, len(body_ids), 3, device=device) + torques[..., 2] = TORQUE_MAGNITUDE + + cube_composer.permanent_wrench_composer.set_forces_and_torques( + forces=forces, + torques=torques, + body_ids=body_ids, + is_global=True, + ) + + # Raw PhysX data + raw_forces = torch.zeros(1, 3, device=device) + raw_torques = torch.zeros(1, 3, device=device) + raw_torques[:, 2] = TORQUE_MAGNITUDE + raw_indices = cube_raw._ALL_INDICES + + for _ in range(N_STEPS): + cube_composer.write_data_to_sim() + cube_raw.write_data_to_sim() + cube_raw.root_view.apply_forces_and_torques_at_position( + force_data=wp.from_torch(raw_forces.contiguous(), dtype=wp.float32), + torque_data=wp.from_torch(raw_torques.contiguous(), dtype=wp.float32), + position_data=None, + indices=raw_indices, + is_global=True, + ) + sim.step() + cube_composer.update(sim.cfg.dt) + cube_raw.update(sim.cfg.dt) + + # Angular velocities should match + torch.testing.assert_close( + wp.to_torch(cube_composer.data.root_ang_vel_w), + wp.to_torch(cube_raw.data.root_ang_vel_w), + rtol=1e-4, + atol=1e-4, + ) + + +NUM_CUBES_MULTI = 4 + + +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +def test_composer_vs_physx_global_force_multi_env(device): + """Global force (no position) with multiple environments. + + Regression: checks that env-indexing and per-body quaternion handling work correctly + when there is more than one environment. + """ + with build_simulation_context(device=device, gravity_enabled=False, auto_add_lighting=True) as sim: + sim._app_control_on_stop_handle = None + cube_composer, cube_raw = generate_dual_cube_scene( + num_cubes=NUM_CUBES_MULTI, device=device, initial_rot=ROT_45_Z + ) + + sim.reset() + + body_ids, _ = cube_composer.find_bodies(".*") + + # Composer path: global force +X for all envs + forces = torch.zeros(NUM_CUBES_MULTI, len(body_ids), 3, device=device) + forces[..., 0] = FORCE_MAGNITUDE + torques = torch.zeros(NUM_CUBES_MULTI, len(body_ids), 3, device=device) + + cube_composer.permanent_wrench_composer.set_forces_and_torques( + forces=forces, + torques=torques, + body_ids=body_ids, + is_global=True, + ) + + # Raw PhysX data (one row per env) + raw_forces = torch.zeros(NUM_CUBES_MULTI, 3, device=device) + raw_forces[:, 0] = FORCE_MAGNITUDE + raw_torques = torch.zeros(NUM_CUBES_MULTI, 3, device=device) + raw_indices = cube_raw._ALL_INDICES + + for _ in range(N_STEPS): + cube_composer.write_data_to_sim() + cube_raw.write_data_to_sim() + cube_raw.root_view.apply_forces_and_torques_at_position( + force_data=wp.from_torch(raw_forces.contiguous(), dtype=wp.float32), + torque_data=wp.from_torch(raw_torques.contiguous(), dtype=wp.float32), + position_data=None, + indices=raw_indices, + is_global=True, + ) + sim.step() + cube_composer.update(sim.cfg.dt) + cube_raw.update(sim.cfg.dt) + + # Linear velocities should match across all envs + torch.testing.assert_close( + wp.to_torch(cube_composer.data.root_lin_vel_w), + wp.to_torch(cube_raw.data.root_lin_vel_w), + rtol=1e-4, + atol=1e-4, + ) + # Angular velocities should match + torch.testing.assert_close( + wp.to_torch(cube_composer.data.root_ang_vel_w), + wp.to_torch(cube_raw.data.root_ang_vel_w), + rtol=1e-4, + atol=1e-4, + ) + # All envs should have ~zero angular velocity + torch.testing.assert_close( + wp.to_torch(cube_composer.data.root_ang_vel_w), + torch.zeros(NUM_CUBES_MULTI, 3, device=device), + rtol=0.0, + atol=1e-4, + ) + + +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +def test_composer_vs_physx_global_force_with_reset(device): + """Global force (no position) with a mid-simulation reset of half the envs. + + Regression: after reset the permanent wrench is cleared. Re-setting it should + produce correct behavior even though the object state was just reset. + """ + with build_simulation_context(device=device, gravity_enabled=False, auto_add_lighting=True) as sim: + sim._app_control_on_stop_handle = None + cube_composer, cube_raw = generate_dual_cube_scene( + num_cubes=NUM_CUBES_MULTI, device=device, initial_rot=ROT_45_Z, spacing=20.0 + ) + + sim.reset() + + # Capture initial world-frame state (includes env origin offsets) + cube_composer.update(sim.cfg.dt) + cube_raw.update(sim.cfg.dt) + initial_state_composer = torch.cat( + [ + wp.to_torch(cube_composer.data.root_link_pos_w), + wp.to_torch(cube_composer.data.root_link_quat_w), + wp.to_torch(cube_composer.data.root_com_vel_w), + ], + dim=-1, + ).clone() + initial_state_raw = torch.cat( + [ + wp.to_torch(cube_raw.data.root_link_pos_w), + wp.to_torch(cube_raw.data.root_link_quat_w), + wp.to_torch(cube_raw.data.root_com_vel_w), + ], + dim=-1, + ).clone() + + body_ids, _ = cube_composer.find_bodies(".*") + + def apply_global_force(): + """Set the same global +X force on the composer cube.""" + forces = torch.zeros(NUM_CUBES_MULTI, len(body_ids), 3, device=device) + forces[..., 0] = FORCE_MAGNITUDE + torques = torch.zeros(NUM_CUBES_MULTI, len(body_ids), 3, device=device) + cube_composer.permanent_wrench_composer.set_forces_and_torques( + forces=forces, + torques=torques, + body_ids=body_ids, + is_global=True, + ) + + apply_global_force() + + # Raw PhysX data + raw_forces = torch.zeros(NUM_CUBES_MULTI, 3, device=device) + raw_forces[:, 0] = FORCE_MAGNITUDE + raw_torques = torch.zeros(NUM_CUBES_MULTI, 3, device=device) + raw_indices = cube_raw._ALL_INDICES + + # Phase 1: run N_STEPS / 2 + half = N_STEPS // 2 + for _ in range(half): + cube_composer.write_data_to_sim() + cube_raw.write_data_to_sim() + cube_raw.root_view.apply_forces_and_torques_at_position( + force_data=wp.from_torch(raw_forces.contiguous(), dtype=wp.float32), + torque_data=wp.from_torch(raw_torques.contiguous(), dtype=wp.float32), + position_data=None, + indices=raw_indices, + is_global=True, + ) + sim.step() + cube_composer.update(sim.cfg.dt) + cube_raw.update(sim.cfg.dt) + + # Reset first half of envs on both cubes + reset_ids = list(range(NUM_CUBES_MULTI // 2)) + reset_ids_torch = torch.tensor(reset_ids, dtype=torch.long, device=device) + + # Reset root state using captured world-frame initial state (includes env origins) + cube_composer.write_root_state_to_sim(initial_state_composer[reset_ids_torch], env_ids=reset_ids_torch) + cube_raw.write_root_state_to_sim(initial_state_raw[reset_ids_torch], env_ids=reset_ids_torch) + + cube_composer.reset(reset_ids) + cube_raw.reset(reset_ids) + + # Re-apply the force (reset cleared the permanent wrench) + apply_global_force() + + # Phase 2: run N_STEPS / 2 more + for _ in range(half): + cube_composer.write_data_to_sim() + cube_raw.write_data_to_sim() + cube_raw.root_view.apply_forces_and_torques_at_position( + force_data=wp.from_torch(raw_forces.contiguous(), dtype=wp.float32), + torque_data=wp.from_torch(raw_torques.contiguous(), dtype=wp.float32), + position_data=None, + indices=raw_indices, + is_global=True, + ) + sim.step() + cube_composer.update(sim.cfg.dt) + cube_raw.update(sim.cfg.dt) + + # All envs: composer vs raw should match + torch.testing.assert_close( + wp.to_torch(cube_composer.data.root_lin_vel_w), + wp.to_torch(cube_raw.data.root_lin_vel_w), + rtol=1e-4, + atol=1e-4, + ) + torch.testing.assert_close( + wp.to_torch(cube_composer.data.root_ang_vel_w), + wp.to_torch(cube_raw.data.root_ang_vel_w), + rtol=1e-4, + atol=1e-4, + ) + # All envs should have ~zero angular velocity + torch.testing.assert_close( + wp.to_torch(cube_composer.data.root_ang_vel_w), + torch.zeros(NUM_CUBES_MULTI, 3, device=device), + rtol=0.0, + atol=1e-4, + ) + + +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +def test_composer_vs_physx_payload_scenario(device): + """Mirrors the apply_payload MDP: permanent global downward force at CoM with gravity. + + A constant world-frame downward force (payload weight) is applied via the composer + path vs raw PhysX. The body falls under gravity + payload, contacts the ground, and + orientation changes. The composer does a world->body->world round-trip each step; + this test catches any precision drift from that. + """ + with build_simulation_context(device=device, gravity_enabled=True, auto_add_lighting=True) as sim: + sim._app_control_on_stop_handle = None + cube_composer, cube_raw = generate_dual_cube_scene( + num_cubes=1, height=0.5, device=device, initial_rot=ROT_45_Z, spacing=20.0 + ) + + sim.reset() + cube_composer.update(sim.cfg.dt) + cube_raw.update(sim.cfg.dt) + + # Record initial positions to compare displacements (cubes spawn at different Y) + init_pos_composer = wp.to_torch(cube_composer.data.root_pos_w).clone() + init_pos_raw = wp.to_torch(cube_raw.data.root_pos_w).clone() + + body_ids, _ = cube_composer.find_bodies(".*") + + payload_force = 2.0 * 9.81 + forces = torch.zeros(1, len(body_ids), 3, device=device) + forces[..., 2] = -payload_force + torques = torch.zeros(1, len(body_ids), 3, device=device) + + cube_composer.permanent_wrench_composer.set_forces_and_torques( + forces=forces, + torques=torques, + body_ids=body_ids, + is_global=True, + ) + + raw_forces = torch.zeros(1, 3, device=device) + raw_forces[:, 2] = -payload_force + raw_torques = torch.zeros(1, 3, device=device) + raw_indices = cube_raw._ALL_INDICES + + for _ in range(N_STEPS): + cube_composer.write_data_to_sim() + cube_raw.write_data_to_sim() + cube_raw.root_view.apply_forces_and_torques_at_position( + force_data=wp.from_torch(raw_forces.contiguous(), dtype=wp.float32), + torque_data=wp.from_torch(raw_torques.contiguous(), dtype=wp.float32), + position_data=None, + indices=raw_indices, + is_global=True, + ) + sim.step() + cube_composer.update(sim.cfg.dt) + cube_raw.update(sim.cfg.dt) + + # Compare displacements (not absolute positions — cubes have different spawn Y) + disp_composer = wp.to_torch(cube_composer.data.root_pos_w) - init_pos_composer + disp_raw = wp.to_torch(cube_raw.data.root_pos_w) - init_pos_raw + + torch.testing.assert_close(disp_composer, disp_raw, rtol=1e-4, atol=1e-4) + torch.testing.assert_close( + wp.to_torch(cube_composer.data.root_lin_vel_w), + wp.to_torch(cube_raw.data.root_lin_vel_w), + rtol=1e-4, + atol=1e-4, + ) + torch.testing.assert_close( + wp.to_torch(cube_composer.data.root_ang_vel_w), + wp.to_torch(cube_raw.data.root_ang_vel_w), + rtol=1e-4, + atol=1e-4, + ) + + +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +def test_composer_vs_physx_permanent_global_force_at_position_long_run(device): + """Permanent global force at a world-frame offset, run long enough for significant body motion. + + This test catches temporal drift bugs where the stored positional torque diverges from + what PhysX computes each step as the body moves. The force is large enough that the body + translates and rotates significantly over 100 steps, but not so large that it causes + numerical instability. + """ + with build_simulation_context(device=device, gravity_enabled=False, auto_add_lighting=True) as sim: + sim._app_control_on_stop_handle = None + cube_composer, cube_raw = generate_dual_cube_scene(num_cubes=1, device=device, initial_rot=ROT_45_Z) + + sim.reset() + + body_ids, _ = cube_composer.find_bodies(".*") + + # Global force +Z at +1m Y offset from CoM — produces torque around X + forces = torch.zeros(1, len(body_ids), 3, device=device) + forces[..., 2] = FORCE_MAGNITUDE + torques = torch.zeros(1, len(body_ids), 3, device=device) + + offset = torch.zeros(1, len(body_ids), 3, device=device) + offset[..., 1] = 1.0 + + pos_composer = wp.to_torch(cube_composer.data.body_com_pos_w)[:, body_ids, :3].clone() + offset + pos_raw = wp.to_torch(cube_raw.data.body_com_pos_w)[:, body_ids, :3].clone() + offset + + cube_composer.permanent_wrench_composer.set_forces_and_torques( + forces=forces, + torques=torques, + positions=pos_composer, + body_ids=body_ids, + is_global=True, + ) + + raw_forces = torch.zeros(1, 3, device=device) + raw_forces[:, 2] = FORCE_MAGNITUDE + raw_torques = torch.zeros(1, 3, device=device) + raw_positions = pos_raw.view(-1, 3) + raw_indices = cube_raw._ALL_INDICES + + for _ in range(100): + cube_composer.write_data_to_sim() + cube_raw.write_data_to_sim() + cube_raw.root_view.apply_forces_and_torques_at_position( + force_data=wp.from_torch(raw_forces.contiguous(), dtype=wp.float32), + torque_data=wp.from_torch(raw_torques.contiguous(), dtype=wp.float32), + position_data=wp.from_torch(raw_positions.contiguous(), dtype=wp.float32), + indices=raw_indices, + is_global=True, + ) + sim.step() + cube_composer.update(sim.cfg.dt) + cube_raw.update(sim.cfg.dt) + + torch.testing.assert_close( + wp.to_torch(cube_composer.data.root_lin_vel_w), + wp.to_torch(cube_raw.data.root_lin_vel_w), + rtol=1e-3, + atol=1e-3, + ) + torch.testing.assert_close( + wp.to_torch(cube_composer.data.root_ang_vel_w), + wp.to_torch(cube_raw.data.root_ang_vel_w), + rtol=1e-3, + atol=1e-3, + ) + + # Sanity: angular velocity should be nonzero + assert torch.abs(wp.to_torch(cube_composer.data.root_ang_vel_w)).max().item() > 0.1, ( + "Expected nonzero angular velocity from positional torque over 100 steps" + ) diff --git a/source/isaaclab/test/visualizers/test_visualizer.py b/source/isaaclab/test/visualizers/test_visualizer.py index ed1baf9198c1..44d3a89aea16 100644 --- a/source/isaaclab/test/visualizers/test_visualizer.py +++ b/source/isaaclab/test/visualizers/test_visualizer.py @@ -7,6 +7,7 @@ from __future__ import annotations +import importlib.util from types import SimpleNamespace import pytest @@ -61,15 +62,18 @@ def is_running(self) -> bool: def _make_cfg(**kwargs): cfg = { - "env_filter_mode": "none", - "env_filter_ids": [0, 2, 4], - "env_filter_random_n": 2, - "env_filter_seed": 7, + "max_visible_envs": None, + "visible_env_indices": None, + # Default off in tests: contiguous cap-only path matches historical assertions. + "randomly_sample_visible_envs": False, } cfg.update(kwargs) return SimpleNamespace(**cfg) +_HAS_ISAACLAB_VIZ = importlib.util.find_spec("isaaclab_visualizers") is not None + + class _FakeProvider: def __init__(self, num_envs: int = 0, transforms: dict | None = None): self._num_envs = num_envs @@ -82,25 +86,68 @@ def get_camera_transforms(self): return self._transforms -def test_compute_visualized_env_ids_none_mode(): - viz = _DummyVisualizer(_make_cfg(env_filter_mode="none")) +def test_compute_visualized_env_ids_cap_only_returns_none(): + """Cap-only path: :meth:`_compute_visualized_env_ids` is ``None``. + + The cap is applied later by ``resolve_visible_env_indices``. + """ + viz = _DummyVisualizer(_make_cfg(visible_env_indices=None)) viz._scene_data_provider = _FakeProvider(num_envs=8) assert viz._compute_visualized_env_ids() is None -def test_compute_visualized_env_ids_from_ids_filters_out_of_range(): - viz = _DummyVisualizer(_make_cfg(env_filter_mode="env_ids", env_filter_ids=[-1, 0, 3, 99])) +def test_compute_visualized_env_ids_from_visible_indices_filters_out_of_range(): + viz = _DummyVisualizer(_make_cfg(visible_env_indices=[-1, 0, 3, 99])) viz._scene_data_provider = _FakeProvider(num_envs=4) assert viz._compute_visualized_env_ids() == [0, 3] -def test_compute_visualized_env_ids_random_n_is_deterministic(): - cfg = _make_cfg(env_filter_mode="random_n", env_filter_random_n=3, env_filter_seed=123) - viz_a = _DummyVisualizer(cfg) - viz_b = _DummyVisualizer(cfg) - viz_a._scene_data_provider = _FakeProvider(num_envs=10) - viz_b._scene_data_provider = _FakeProvider(num_envs=10) - assert viz_a._compute_visualized_env_ids() == viz_b._compute_visualized_env_ids() +@pytest.mark.skipif(not _HAS_ISAACLAB_VIZ, reason="isaaclab_visualizers not installed") +def test_partial_visualization_cap_only_uses_resolver(): + """With ``visible_env_indices`` unset, :func:`resolve_visible_env_indices` applies ``max_visible_envs``.""" + from isaaclab_visualizers.newton_adapter import resolve_visible_env_indices + + cfg = _make_cfg(max_visible_envs=3, visible_env_indices=None) + viz = _DummyVisualizer(cfg) + viz._scene_data_provider = _FakeProvider(num_envs=10) + assert viz._compute_visualized_env_ids() is None + assert resolve_visible_env_indices(None, cfg.max_visible_envs, 10) == [0, 1, 2] + assert resolve_visible_env_indices(None, 3, 10) == [0, 1, 2] + + +@pytest.mark.skipif(not _HAS_ISAACLAB_VIZ, reason="isaaclab_visualizers not installed") +def test_compute_visualized_env_ids_random_cap_only_sorted_once(): + """Cap-only random mode returns a sorted sample; explicit indices ignore the flag.""" + cfg = _make_cfg(max_visible_envs=3, visible_env_indices=None, randomly_sample_visible_envs=True) + viz = _DummyVisualizer(cfg) + viz._scene_data_provider = _FakeProvider(num_envs=10) + sampled = viz._compute_visualized_env_ids() + assert sampled is not None and len(sampled) == 3 + assert sampled == sorted(sampled) + assert len(set(sampled)) == 3 + assert all(0 <= i < 10 for i in sampled) + + cfg_explicit = _make_cfg( + visible_env_indices=[1, 5], + max_visible_envs=1, + randomly_sample_visible_envs=True, + ) + viz2 = _DummyVisualizer(cfg_explicit) + viz2._scene_data_provider = _FakeProvider(num_envs=10) + assert viz2._compute_visualized_env_ids() == [1, 5] + + +@pytest.mark.skipif(not _HAS_ISAACLAB_VIZ, reason="isaaclab_visualizers not installed") +def test_explicit_visible_env_indices_truncated_by_max_visible_envs(): + """Explicit indices from :meth:`_compute_visualized_env_ids`; ``max_visible_envs`` truncates from the end.""" + from isaaclab_visualizers.newton_adapter import resolve_visible_env_indices + + cfg = _make_cfg(visible_env_indices=[0, 2, 4], max_visible_envs=1) + viz = _DummyVisualizer(cfg) + viz._scene_data_provider = _FakeProvider(num_envs=10) + ids = viz._compute_visualized_env_ids() + assert ids == [0, 2, 4] + assert resolve_visible_env_indices(ids, cfg.max_visible_envs, 10) == [0] def test_resolve_camera_pose_from_usd_path_uses_provider_transforms(): 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 fe15054f69cd..7d5691efdeb4 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.17" +version = "0.5.21" # 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 0230a394935e..ec1f45bec732 100644 --- a/source/isaaclab_newton/docs/CHANGELOG.rst +++ b/source/isaaclab_newton/docs/CHANGELOG.rst @@ -1,6 +1,56 @@ Changelog --------- +0.5.21 (2026-04-23) +~~~~~~~~~~~~~~~~~~~ + +Fixed +^^^^^ + +* Fixed flakiness in ``test_body_root_state_properties`` by bounding the random spin velocity so + numerical drift stays within the position tolerance over the simulated trajectory. + + +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) +~~~~~~~~~~~~~~~~~~~ + +Changed +^^^^^^^ + +* Updated ``write_data_to_sim`` in :class:`~isaaclab_newton.assets.Articulation`, + :class:`~isaaclab_newton.assets.RigidObject`, and :class:`~isaaclab_newton.assets.RigidObjectCollection` + to use the dual-buffer :class:`~isaaclab.utils.wrench_composer.WrenchComposer`. Composed wrenches are + applied after body-frame composition. + + +0.5.18 (2026-04-21) +~~~~~~~~~~~~~~~~~~~ + +Changed +^^^^^^^ + +* Upgraded Newton from ``2684d75`` to ``a27277e``. Includes collision improvements, contact quality fixes, + hydroelastic contact optimization, and memory usage fixes in CollisionPipeline. For details see + ``Newton changelog ``. +* Pinned ``mujoco`` and ``mujoco-warp`` to ``3.6.0`` to align with the Newton library. + + 0.5.17 (2026-04-20) ~~~~~~~~~~~~~~~~~~~ @@ -105,10 +155,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/assets/articulation/articulation.py b/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py index 9d62dc0bbed1..515176352490 100644 --- a/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py +++ b/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py @@ -256,40 +256,23 @@ def write_data_to_sim(self): # write external wrench if self._instantaneous_wrench_composer.active or self._permanent_wrench_composer.active: if self._instantaneous_wrench_composer.active: - # Compose instantaneous wrench with permanent wrench - self._instantaneous_wrench_composer.add_forces_and_torques_index( - forces=self._permanent_wrench_composer.composed_force, - torques=self._permanent_wrench_composer.composed_torque, - body_ids=self._ALL_BODY_INDICES, - env_ids=self._ALL_INDICES, - ) - # Apply both instantaneous and permanent wrench to the simulation - wp.launch( - shared_kernels.update_wrench_array_with_force_and_torque, - dim=(self.num_instances, self.num_bodies), - device=self.device, - inputs=[ - self._instantaneous_wrench_composer.composed_force, - self._instantaneous_wrench_composer.composed_torque, - self._data._sim_bind_body_external_wrench, - self._ALL_ENV_MASK, - self._ALL_BODY_MASK, - ], - ) + composer = self._instantaneous_wrench_composer + composer.add_raw_buffers_from(self._permanent_wrench_composer) else: - # Apply permanent wrench to the simulation - wp.launch( - shared_kernels.update_wrench_array_with_force_and_torque, - dim=(self.num_instances, self.num_bodies), - device=self.device, - inputs=[ - self._permanent_wrench_composer.composed_force, - self._permanent_wrench_composer.composed_torque, - self._data._sim_bind_body_external_wrench, - self._ALL_ENV_MASK, - self._ALL_BODY_MASK, - ], - ) + composer = self._permanent_wrench_composer + composer.compose_to_body_frame() + wp.launch( + shared_kernels.update_wrench_array_with_force_and_torque, + dim=(self.num_instances, self.num_bodies), + device=self.device, + inputs=[ + composer.out_force_b, + composer.out_torque_b, + self._data._sim_bind_body_external_wrench, + self._ALL_ENV_MASK, + self._ALL_BODY_MASK, + ], + ) self._instantaneous_wrench_composer.reset() # apply actuator models diff --git a/source/isaaclab_newton/isaaclab_newton/assets/rigid_object/rigid_object.py b/source/isaaclab_newton/isaaclab_newton/assets/rigid_object/rigid_object.py index 5e02e3622985..fb2d29091203 100644 --- a/source/isaaclab_newton/isaaclab_newton/assets/rigid_object/rigid_object.py +++ b/source/isaaclab_newton/isaaclab_newton/assets/rigid_object/rigid_object.py @@ -143,40 +143,23 @@ def write_data_to_sim(self) -> None: # write external wrench if self._instantaneous_wrench_composer.active or self._permanent_wrench_composer.active: if self._instantaneous_wrench_composer.active: - # Compose instantaneous wrench with permanent wrench - self._instantaneous_wrench_composer.add_forces_and_torques_index( - forces=self._permanent_wrench_composer.composed_force, - torques=self._permanent_wrench_composer.composed_torque, - body_ids=self._ALL_BODY_INDICES, - env_ids=self._ALL_INDICES, - ) - # Apply both instantaneous and permanent wrench to the simulation - wp.launch( - shared_kernels.update_wrench_array_with_force_and_torque, - dim=(self.num_instances, self.num_bodies), - device=self.device, - inputs=[ - self._instantaneous_wrench_composer.composed_force, - self._instantaneous_wrench_composer.composed_torque, - self._data._sim_bind_body_external_wrench, - self._ALL_ENV_MASK, - self._ALL_BODY_MASK, - ], - ) + composer = self._instantaneous_wrench_composer + composer.add_raw_buffers_from(self._permanent_wrench_composer) else: - # Apply permanent wrench to the simulation - wp.launch( - shared_kernels.update_wrench_array_with_force_and_torque, - dim=(self.num_instances, self.num_bodies), - device=self.device, - inputs=[ - self._permanent_wrench_composer.composed_force, - self._permanent_wrench_composer.composed_torque, - self._data._sim_bind_body_external_wrench, - self._ALL_ENV_MASK, - self._ALL_BODY_MASK, - ], - ) + composer = self._permanent_wrench_composer + composer.compose_to_body_frame() + wp.launch( + shared_kernels.update_wrench_array_with_force_and_torque, + dim=(self.num_instances, self.num_bodies), + device=self.device, + inputs=[ + composer.out_force_b, + composer.out_torque_b, + self._data._sim_bind_body_external_wrench, + self._ALL_ENV_MASK, + self._ALL_BODY_MASK, + ], + ) self._instantaneous_wrench_composer.reset() def update(self, dt: float) -> None: diff --git a/source/isaaclab_newton/isaaclab_newton/assets/rigid_object_collection/rigid_object_collection.py b/source/isaaclab_newton/isaaclab_newton/assets/rigid_object_collection/rigid_object_collection.py index 82bdf7a03003..52216290d329 100644 --- a/source/isaaclab_newton/isaaclab_newton/assets/rigid_object_collection/rigid_object_collection.py +++ b/source/isaaclab_newton/isaaclab_newton/assets/rigid_object_collection/rigid_object_collection.py @@ -187,40 +187,23 @@ def write_data_to_sim(self) -> None: # write external wrench if self._instantaneous_wrench_composer.active or self._permanent_wrench_composer.active: if self._instantaneous_wrench_composer.active: - # Compose instantaneous wrench with permanent wrench - self._instantaneous_wrench_composer.add_forces_and_torques_index( - forces=self._permanent_wrench_composer.composed_force, - torques=self._permanent_wrench_composer.composed_torque, - body_ids=self._ALL_BODY_INDICES, - env_ids=self._ALL_ENV_INDICES, - ) - # Apply both instantaneous and permanent wrench to a consolidated 2D buffer - wp.launch( - shared_kernels.update_wrench_array_with_force_and_torque, - dim=(self.num_instances, self.num_bodies), - device=self.device, - inputs=[ - self._instantaneous_wrench_composer.composed_force, - self._instantaneous_wrench_composer.composed_torque, - self._wrench_buffer, - self._ALL_ENV_MASK, - self._ALL_BODY_MASK, - ], - ) + composer = self._instantaneous_wrench_composer + composer.add_raw_buffers_from(self._permanent_wrench_composer) else: - # Apply permanent wrench to a consolidated 2D buffer - wp.launch( - shared_kernels.update_wrench_array_with_force_and_torque, - dim=(self.num_instances, self.num_bodies), - device=self.device, - inputs=[ - self._permanent_wrench_composer.composed_force, - self._permanent_wrench_composer.composed_torque, - self._wrench_buffer, - self._ALL_ENV_MASK, - self._ALL_BODY_MASK, - ], - ) + composer = self._permanent_wrench_composer + composer.compose_to_body_frame() + wp.launch( + shared_kernels.update_wrench_array_with_force_and_torque, + dim=(self.num_instances, self.num_bodies), + device=self.device, + inputs=[ + composer.out_force_b, + composer.out_torque_b, + self._wrench_buffer, + self._ALL_ENV_MASK, + self._ALL_BODY_MASK, + ], + ) # Write the wrench buffer directly to the Newton binding (already 2D) wp.copy(self._data._sim_bind_body_external_wrench, self._wrench_buffer) self._instantaneous_wrench_composer.reset() diff --git a/source/isaaclab_newton/isaaclab_newton/renderers/newton_warp_renderer.py b/source/isaaclab_newton/isaaclab_newton/renderers/newton_warp_renderer.py index a870f0ae0530..22f8acb47a93 100644 --- a/source/isaaclab_newton/isaaclab_newton/renderers/newton_warp_renderer.py +++ b/source/isaaclab_newton/isaaclab_newton/renderers/newton_warp_renderer.py @@ -7,7 +7,6 @@ from __future__ import annotations -import functools import logging import weakref from dataclasses import dataclass @@ -30,36 +29,6 @@ logger = logging.getLogger(__name__) -try: - import nvtx - - _nvtx_domain = nvtx.Domain("NewtonWarpRenderer") - - def _nvtx_range(message: str, color: str | None = None): - """Decorator that wraps a function in a Domain.push_range/pop_range pair.""" - attrs = _nvtx_domain.get_event_attributes(message=message, color=color) - - def decorator(fn): - @functools.wraps(fn) - def wrapper(*args, **kwargs): - _nvtx_domain.push_range(attrs) - try: - return fn(*args, **kwargs) - finally: - _nvtx_domain.pop_range() - - return wrapper - - return decorator - -except ImportError: - - def _nvtx_range(message: str, color: str | None = None): - def decorator(fn): - return fn - - return decorator - class RenderData: class OutputNames: @@ -230,7 +199,6 @@ def set_outputs(self, render_data: RenderData, output_data: dict[str, torch.Tens """Store output buffers. See :meth:`~isaaclab.renderers.base_renderer.BaseRenderer.set_outputs`.""" render_data.set_outputs(output_data) - @_nvtx_range("update_transforms", color="blue") def update_transforms(self): """Sync Newton scene state before rendering. See :meth:`~isaaclab.renderers.base_renderer.BaseRenderer.update_transforms`.""" @@ -243,7 +211,6 @@ def update_camera( See :meth:`~isaaclab.renderers.base_renderer.BaseRenderer.update_camera`.""" render_data.update(positions, orientations, intrinsics) - @_nvtx_range("render", color="green") def render(self, render_data: RenderData): """Render and write to output buffers. See :meth:`~isaaclab.renderers.base_renderer.BaseRenderer.render`.""" self.newton_sensor.update( @@ -259,7 +226,6 @@ def render(self, render_data: RenderData): clear_data=newton.sensors.SensorTiledCamera.ClearData(clear_color=0xFFEEEEEE), ) - @_nvtx_range("read_output", color="orange") def read_output(self, render_data: RenderData, camera_data: CameraData) -> None: """Copy rendered outputs to the camera data buffers. See :meth:`~isaaclab.renderers.base_renderer.BaseRenderer.read_output`.""" diff --git a/source/isaaclab_newton/isaaclab_newton/scene_data_providers/newton_scene_data_provider.py b/source/isaaclab_newton/isaaclab_newton/scene_data_providers/newton_scene_data_provider.py index f3b10dc40044..ba19f4e7c63a 100644 --- a/source/isaaclab_newton/isaaclab_newton/scene_data_providers/newton_scene_data_provider.py +++ b/source/isaaclab_newton/isaaclab_newton/scene_data_providers/newton_scene_data_provider.py @@ -96,7 +96,7 @@ def _determine_num_envs_in_scene(self) -> int: # ---- Core provider API ------------------------------------------------------------------- - def update(self, env_ids: list[int] | None = None) -> None: + def update(self) -> None: """Sync Newton body transforms to USD Fabric when a Kit viewport is active. Called at render cadence by :meth:`~isaaclab.sim.SimulationContext.update_scene_data_provider`, @@ -104,9 +104,6 @@ def update(self, env_ids: list[int] | None = None) -> None: :meth:`~isaaclab_newton.physics.NewtonManager.sync_transforms_to_usd` when a Kit (or other USD-based) visualizer is in use. When both sim and rendering backend are Newton (or Rerun), the sync is skipped to avoid unnecessary slowdown. - - Args: - env_ids: Optional environment id selection. Unused in this provider. """ if not self._needs_usd_sync: return @@ -127,13 +124,9 @@ def get_newton_model(self) -> Any | None: return NewtonManager.get_model() - def get_newton_state(self, env_ids: list[int] | None = None) -> Any | None: + def get_newton_state(self) -> Any | None: """Return Newton state from NewtonManager. - Args: - env_ids: Optional list of environment IDs. Currently returns the full - state for all environments (env_ids filtering is not yet implemented). - Returns: The current Newton state (state_0) from NewtonManager. """ @@ -149,16 +142,9 @@ def get_model(self) -> Any | None: """ return self.get_newton_model() - def get_state(self, env_ids: list[int] | None = None) -> Any | None: - """Alias for :meth:`get_newton_state` for visualizer compatibility. - - Args: - env_ids: Optional list of environment ids. - - Returns: - Newton state object, or ``None`` when unavailable. - """ - return self.get_newton_state(env_ids) + def get_state(self) -> Any | None: + """Alias for :meth:`get_newton_state` for visualizer compatibility.""" + return self.get_newton_state() def get_usd_stage(self) -> Any | None: """Return the USD stage handle. 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/isaaclab_newton/video_recording/__init__.py b/source/isaaclab_newton/isaaclab_newton/video_recording/__init__.py index 3248ca5f13b4..1d5cb96e0ef3 100644 --- a/source/isaaclab_newton/isaaclab_newton/video_recording/__init__.py +++ b/source/isaaclab_newton/isaaclab_newton/video_recording/__init__.py @@ -4,3 +4,8 @@ # SPDX-License-Identifier: BSD-3-Clause """Newton GL perspective video recording.""" + +from .newton_gl_perspective_video import NewtonGlPerspectiveVideo +from .newton_gl_perspective_video_cfg import NewtonGlPerspectiveVideoCfg + +__all__ = ["NewtonGlPerspectiveVideo", "NewtonGlPerspectiveVideoCfg"] diff --git a/source/isaaclab_newton/isaaclab_newton/video_recording/recording_hooks.py b/source/isaaclab_newton/isaaclab_newton/video_recording/recording_hooks.py new file mode 100644 index 000000000000..7efcae7b5500 --- /dev/null +++ b/source/isaaclab_newton/isaaclab_newton/video_recording/recording_hooks.py @@ -0,0 +1,30 @@ +# 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 + +"""Hooks for Newton-based video recording after visualizers have stepped.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from isaaclab.sim import SimulationContext + + +def recording_followup_after_visualizers(sim: SimulationContext) -> None: + """Newton extension hook: recording pipeline after visualizers have stepped. + + Called from :func:`isaaclab.envs.utils.recording_hooks.run_recording_hooks_after_visualizers`. + Wire **Newton GL** / Newton-specific video capture here (e.g. perspective video, + frame sync with ``NewtonVisualizer``). Stay lightweight and no-op when Newton + recording is inactive. + + The Isaac Sim / RTX path (``omni.kit.app`` pump for Replicator ``rgb_array``) lives in + :mod:`isaaclab_physx.renderers.isaac_rtx_renderer_utils` — not here. + + Args: + sim: Active simulation context. + """ + _ = sim # Reserved until Newton GL video paths are hooked up. diff --git a/source/isaaclab_newton/setup.py b/source/isaaclab_newton/setup.py index 421cecd502ca..2e0b87f17543 100644 --- a/source/isaaclab_newton/setup.py +++ b/source/isaaclab_newton/setup.py @@ -33,19 +33,15 @@ def run(self): # Read the extension.toml file EXTENSION_TOML_DATA = toml.load(os.path.join(EXTENSION_PATH, "config", "extension.toml")) -INSTALL_REQUIRES = [ - # INTENTIONALLY disabled to avoid circular dependency with isaaclab_physx, which also depends on isaaclab_newton. - # This will be re-enabled once we move to UV and pyproject.toml-based packaging. - # f"isaaclab_physx @ file://{os.path.join(os.path.dirname(EXTENSION_PATH), 'isaaclab_physx')}", -] +INSTALL_REQUIRES = [] EXTRAS_REQUIRE = { "all": [ "prettytable==3.3.0", - "mujoco==3.5.0", - "mujoco-warp==3.5.0.2", + "mujoco==3.6.0", + "mujoco-warp==3.6.0", "PyOpenGL-accelerate==3.1.10", - "newton @ git+https://github.com/newton-physics/newton.git@2684d75bfa4bb8b058a93b81c458a74b7701c997", + "newton @ git+https://github.com/newton-physics/newton.git@a27277ed49d6f307b8a1e4c394be7e1d14965a62", ], } diff --git a/source/isaaclab_newton/test/assets/test_rigid_object.py b/source/isaaclab_newton/test/assets/test_rigid_object.py index c1d01f4164fb..152f55d8c6f4 100644 --- a/source/isaaclab_newton/test/assets/test_rigid_object.py +++ b/source/isaaclab_newton/test/assets/test_rigid_object.py @@ -372,7 +372,6 @@ def test_external_force_on_single_body(num_cubes, device): assert torch.all(wp.to_torch(cube_object.data.root_pos_w)[1::2, 2] < 1.0) -@pytest.mark.skip(reason="Newton wrench composer at-position force composition differs from PhysX") @pytest.mark.parametrize("num_cubes", [2, 4]) @pytest.mark.parametrize("device", ["cuda:0", "cpu"]) def test_external_force_on_single_body_at_position(num_cubes, device): @@ -399,14 +398,9 @@ def test_external_force_on_single_body_at_position(num_cubes, device): external_wrench_b = torch.zeros(cube_object.num_instances, len(body_ids), 6, device=sim.device) external_wrench_positions_b = torch.zeros(cube_object.num_instances, len(body_ids), 3, device=sim.device) # Every 2nd cube should have a force applied to it - external_wrench_b[0::2, :, 2] = 500.0 + external_wrench_b[0::2, :, 2] = 50.0 external_wrench_positions_b[0::2, :, 1] = 1.0 - # Desired force and torque - desired_force = torch.zeros(cube_object.num_instances, len(body_ids), 3, device=sim.device) - desired_force[0::2, :, 2] = 1000.0 - desired_torque = torch.zeros(cube_object.num_instances, len(body_ids), 3, device=sim.device) - desired_torque[0::2, :, 0] = 1000.0 # Now we are ready! for i in range(5): # reset root state @@ -449,18 +443,6 @@ def test_external_force_on_single_body_at_position(num_cubes, device): body_ids=body_ids, is_global=is_global, ) - torch.testing.assert_close( - wp.to_torch(cube_object._permanent_wrench_composer.composed_force)[:, 0, :], - desired_force[:, 0, :], - rtol=1e-6, - atol=1e-7, - ) - torch.testing.assert_close( - wp.to_torch(cube_object._permanent_wrench_composer.composed_torque)[:, 0, :], - desired_torque[:, 0, :], - rtol=1e-6, - atol=1e-7, - ) # perform simulation for _ in range(5): # apply action to the object @@ -966,9 +948,9 @@ def test_body_root_state_properties(num_cubes, device, with_offset): # check center of mass has been set torch.testing.assert_close(wp.to_torch(cube_object.data.body_com_pos_b).squeeze(1), offset) - # random z spin velocity + # random z spin velocity (bounded to keep numerical drift within the position tolerance below) spin_twist = torch.zeros(6, device=device) - spin_twist[5] = torch.randn(1, device=device) + spin_twist[5] = 0.5 * torch.randn(1, device=device).clamp(-1.0, 1.0) # Simulate physics for _ in range(100): diff --git a/source/isaaclab_newton/test/assets/test_rigid_object_collection.py b/source/isaaclab_newton/test/assets/test_rigid_object_collection.py index 7d4a0be7cb98..4c5599e35887 100644 --- a/source/isaaclab_newton/test/assets/test_rigid_object_collection.py +++ b/source/isaaclab_newton/test/assets/test_rigid_object_collection.py @@ -344,7 +344,7 @@ def test_external_force_on_single_body_at_position(num_envs, num_cubes, device): object_collection.num_instances, len(object_ids), 3, device=sim.device ) # Every 2nd cube should have a force applied to it - external_wrench_b[:, 0::2, 2] = 500.0 + external_wrench_b[:, 0::2, 2] = 50.0 external_wrench_positions_b[:, 0::2, 1] = 1.0 # Desired force and torque 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_ovphysx/config/extension.toml b/source/isaaclab_ovphysx/config/extension.toml index 11b3322f2f1d..ed4f5b39fb70 100644 --- a/source/isaaclab_ovphysx/config/extension.toml +++ b/source/isaaclab_ovphysx/config/extension.toml @@ -1,7 +1,7 @@ [package] # Note: Semantic Versioning is used: https://semver.org/ -version = "0.1.0" +version = "0.1.1" # Description title = "OvPhysX simulation interfaces for IsaacLab core package" diff --git a/source/isaaclab_ovphysx/docs/CHANGELOG.rst b/source/isaaclab_ovphysx/docs/CHANGELOG.rst index b177752442d3..750a0397f23d 100644 --- a/source/isaaclab_ovphysx/docs/CHANGELOG.rst +++ b/source/isaaclab_ovphysx/docs/CHANGELOG.rst @@ -1,6 +1,24 @@ Changelog --------- +0.1.1 (2026-04-21) +~~~~~~~~~~~~~~~~~~~ + +Changed +^^^^^^^ + +* Replaced private ``_find_names`` (fnmatch + regex) with the standard + :func:`~isaaclab.utils.string.resolve_matching_names` for all finder + methods, unifying name-resolution behavior across backends. Fnmatch-style + glob patterns (e.g. ``joint_*``) are no longer supported; use regex + equivalents (e.g. ``joint_.*``). ``find_fixed_tendons`` and + ``find_spatial_tendons`` now raise ``ValueError`` on empty tendon lists, + matching the PhysX backend. +* Changed ``find_joints`` ``joint_subset`` parameter from ``list[int]`` + (indices) to ``list[str]`` (names) to match the ``BaseArticulation`` + interface. Callers passing indices should convert to names first. + + 0.1.0 (2026-04-20) ~~~~~~~~~~~~~~~~~~ diff --git a/source/isaaclab_ovphysx/isaaclab_ovphysx/assets/articulation/articulation.py b/source/isaaclab_ovphysx/isaaclab_ovphysx/assets/articulation/articulation.py index 4c00dc839ca1..7224d53d40ea 100644 --- a/source/isaaclab_ovphysx/isaaclab_ovphysx/assets/articulation/articulation.py +++ b/source/isaaclab_ovphysx/isaaclab_ovphysx/assets/articulation/articulation.py @@ -7,7 +7,6 @@ from __future__ import annotations -import fnmatch import logging import re from collections.abc import Sequence @@ -19,6 +18,7 @@ from isaaclab.assets.articulation.base_articulation import BaseArticulation from isaaclab.physics import PhysicsManager +from isaaclab.utils.string import resolve_matching_names from isaaclab.utils.wrench_composer import WrenchComposer from isaaclab_ovphysx import tensor_types as TT @@ -187,12 +187,12 @@ def find_bodies(self, name_keys: str | Sequence[str], preserve_order: bool = Fal Returns: A tuple of lists containing the body indices and names. """ - return self._find_names(self._body_names, name_keys, preserve_order) + return resolve_matching_names(name_keys, self._body_names, preserve_order) def find_joints( self, name_keys: str | Sequence[str], - joint_subset: list[int] | None = None, + joint_subset: list[str] | None = None, preserve_order: bool = False, ) -> tuple[list[int], list[str]]: """Find joints in the articulation based on the name keys. @@ -202,18 +202,16 @@ def find_joints( Args: name_keys: A regular expression or a list of regular expressions to match the joint names. - joint_subset: A subset of joint indices to search within. Defaults to None, which means all joints + joint_subset: A subset of joints to search for. Defaults to None, which means all joints in the articulation are searched. preserve_order: Whether to preserve the order of the name keys in the output. Defaults to False. Returns: A tuple of lists containing the joint indices and names. """ - names = [self._joint_names[i] for i in joint_subset] if joint_subset is not None else self._joint_names - indices, matched = self._find_names(names, name_keys, preserve_order) - if joint_subset is not None: - indices = [joint_subset[i] for i in indices] - return indices, matched + if joint_subset is None: + joint_subset = self._joint_names + return resolve_matching_names(name_keys, joint_subset, preserve_order) def find_fixed_tendons( self, @@ -237,9 +235,7 @@ def find_fixed_tendons( """ if tendon_subsets is None: tendon_subsets = self.fixed_tendon_names - if not tendon_subsets: - return [], [] - return self._find_names(tendon_subsets, name_keys, preserve_order) + return resolve_matching_names(name_keys, tendon_subsets, preserve_order) def find_spatial_tendons( self, @@ -262,9 +258,7 @@ def find_spatial_tendons( """ if tendon_subsets is None: tendon_subsets = self.spatial_tendon_names - if not tendon_subsets: - return [], [] - return self._find_names(tendon_subsets, name_keys, preserve_order) + return resolve_matching_names(name_keys, tendon_subsets, preserve_order) """ Operations - State Writers. @@ -2664,28 +2658,6 @@ def _nst(self): """Return the number of spatial tendons (0 if none).""" return getattr(self, "_num_spatial_tendons", 0) - @staticmethod - def _find_names(names: list[str], keys: str | Sequence[str], preserve_order: bool) -> tuple[list[int], list[str]]: - if isinstance(keys, str): - keys = [keys] - matched_indices: list[int] = [] - matched_names: list[str] = [] - if preserve_order: - for key in keys: - for idx, name in enumerate(names): - if fnmatch.fnmatch(name, key) or re.fullmatch(key, name): - if idx not in matched_indices: - matched_indices.append(idx) - matched_names.append(name) - else: - for idx, name in enumerate(names): - for key in keys: - if fnmatch.fnmatch(name, key) or re.fullmatch(key, name): - matched_indices.append(idx) - matched_names.append(name) - break - return matched_indices, matched_names - def _resolve_joint_values(self, pattern_dict: dict[str, float], buffer: wp.array) -> None: """Resolve a {pattern: value} dict into a per-joint buffer. diff --git a/source/isaaclab_physx/config/extension.toml b/source/isaaclab_physx/config/extension.toml index 555928ce4c23..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.19" +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 99432eea87a3..a343bf0367fd 100644 --- a/source/isaaclab_physx/docs/CHANGELOG.rst +++ b/source/isaaclab_physx/docs/CHANGELOG.rst @@ -1,6 +1,34 @@ 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) +~~~~~~~~~~~~~~~~~~~ + +Changed +^^^^^^^ + +* Updated ``write_data_to_sim`` in :class:`~isaaclab_physx.assets.Articulation`, + :class:`~isaaclab_physx.assets.RigidObject`, and :class:`~isaaclab_physx.assets.RigidObjectCollection` + to use the dual-buffer :class:`~isaaclab.utils.wrench_composer.WrenchComposer`. Composed wrenches are + applied to PhysX with ``is_global=False`` after body-frame composition. + + 0.5.19 (2026-04-20) ~~~~~~~~~~~~~~~~~~~ diff --git a/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py b/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py index cf7d1f95d5ca..3b403ee8c6d4 100644 --- a/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py +++ b/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py @@ -237,30 +237,18 @@ def write_data_to_sim(self): # write external wrench if self._instantaneous_wrench_composer.active or self._permanent_wrench_composer.active: if self._instantaneous_wrench_composer.active: - # Compose instantaneous wrench with permanent wrench - self._instantaneous_wrench_composer.add_forces_and_torques_index( - forces=self._permanent_wrench_composer.composed_force, - torques=self._permanent_wrench_composer.composed_torque, - body_ids=self._ALL_BODY_INDICES, - env_ids=self._ALL_INDICES, - ) - # Apply both instantaneous and permanent wrench to the simulation - self.root_view.apply_forces_and_torques_at_position( - force_data=self._instantaneous_wrench_composer.composed_force.flatten().view(wp.float32), - torque_data=self._instantaneous_wrench_composer.composed_torque.flatten().view(wp.float32), - position_data=None, - indices=self._ALL_INDICES, - is_global=False, - ) + composer = self._instantaneous_wrench_composer + composer.add_raw_buffers_from(self._permanent_wrench_composer) else: - # Apply permanent wrench to the simulation - self.root_view.apply_forces_and_torques_at_position( - force_data=self._permanent_wrench_composer.composed_force.flatten().view(wp.float32), - torque_data=self._permanent_wrench_composer.composed_torque.flatten().view(wp.float32), - position_data=None, - indices=self._ALL_INDICES, - is_global=False, - ) + composer = self._permanent_wrench_composer + composer.compose_to_body_frame() + self.root_view.apply_forces_and_torques_at_position( + force_data=composer.out_force_b.flatten().view(wp.float32), + torque_data=composer.out_torque_b.flatten().view(wp.float32), + position_data=None, + indices=self._ALL_INDICES, + is_global=False, + ) self._instantaneous_wrench_composer.reset() # apply actuator models diff --git a/source/isaaclab_physx/isaaclab_physx/assets/rigid_object/rigid_object.py b/source/isaaclab_physx/isaaclab_physx/assets/rigid_object/rigid_object.py index b549da96b787..8aa7dbd3f4f3 100644 --- a/source/isaaclab_physx/isaaclab_physx/assets/rigid_object/rigid_object.py +++ b/source/isaaclab_physx/isaaclab_physx/assets/rigid_object/rigid_object.py @@ -150,30 +150,18 @@ def write_data_to_sim(self) -> None: # write external wrench if self._instantaneous_wrench_composer.active or self._permanent_wrench_composer.active: if self._instantaneous_wrench_composer.active: - # Compose instantaneous wrench with permanent wrench - self._instantaneous_wrench_composer.add_forces_and_torques_index( - forces=self._permanent_wrench_composer.composed_force, - torques=self._permanent_wrench_composer.composed_torque, - body_ids=self._ALL_BODY_INDICES, - env_ids=self._ALL_INDICES, - ) - # Apply both instantaneous and permanent wrench to the simulation - self.root_view.apply_forces_and_torques_at_position( - force_data=self._instantaneous_wrench_composer.composed_force.flatten().view(wp.float32), - torque_data=self._instantaneous_wrench_composer.composed_torque.flatten().view(wp.float32), - position_data=None, - indices=self._ALL_INDICES, - is_global=False, - ) + composer = self._instantaneous_wrench_composer + composer.add_raw_buffers_from(self._permanent_wrench_composer) else: - # Apply permanent wrench to the simulation - self.root_view.apply_forces_and_torques_at_position( - force_data=self._permanent_wrench_composer.composed_force.flatten().view(wp.float32), - torque_data=self._permanent_wrench_composer.composed_torque.flatten().view(wp.float32), - position_data=None, - indices=self._ALL_INDICES, - is_global=False, - ) + composer = self._permanent_wrench_composer + composer.compose_to_body_frame() + self.root_view.apply_forces_and_torques_at_position( + force_data=composer.out_force_b.flatten().view(wp.float32), + torque_data=composer.out_torque_b.flatten().view(wp.float32), + position_data=None, + indices=self._ALL_INDICES, + is_global=False, + ) self._instantaneous_wrench_composer.reset() def update(self, dt: float) -> None: diff --git a/source/isaaclab_physx/isaaclab_physx/assets/rigid_object_collection/rigid_object_collection.py b/source/isaaclab_physx/isaaclab_physx/assets/rigid_object_collection/rigid_object_collection.py index 877a44261293..3518aceac1d9 100644 --- a/source/isaaclab_physx/isaaclab_physx/assets/rigid_object_collection/rigid_object_collection.py +++ b/source/isaaclab_physx/isaaclab_physx/assets/rigid_object_collection/rigid_object_collection.py @@ -186,42 +186,20 @@ def write_data_to_sim(self) -> None: # write external wrench if self._instantaneous_wrench_composer.active or self._permanent_wrench_composer.active: if self._instantaneous_wrench_composer.active: - # Compose instantaneous wrench with permanent wrench - self._instantaneous_wrench_composer.add_forces_and_torques_index( - forces=self._permanent_wrench_composer.composed_force, - torques=self._permanent_wrench_composer.composed_torque, - body_ids=self._ALL_BODY_INDICES, - env_ids=self._ALL_ENV_INDICES, - ) - # Apply both instantaneous and permanent wrench to the simulation - self.root_view.apply_forces_and_torques_at_position( - force_data=self.reshape_data_to_view_2d( - self._instantaneous_wrench_composer.composed_force, device=self.device - ).view(wp.float32), - torque_data=self.reshape_data_to_view_2d( - self._instantaneous_wrench_composer.composed_torque, device=self.device - ).view(wp.float32), - position_data=None, - indices=self._env_body_ids_to_view_ids( - self._ALL_ENV_INDICES, self._ALL_BODY_INDICES, device=self.device - ), - is_global=False, - ) + composer = self._instantaneous_wrench_composer + composer.add_raw_buffers_from(self._permanent_wrench_composer) else: - # Apply permanent wrench to the simulation - self.root_view.apply_forces_and_torques_at_position( - force_data=self.reshape_data_to_view_2d( - self._permanent_wrench_composer.composed_force, device=self.device - ).view(wp.float32), - torque_data=self.reshape_data_to_view_2d( - self._permanent_wrench_composer.composed_torque, device=self.device - ).view(wp.float32), - position_data=None, - indices=self._env_body_ids_to_view_ids( - self._ALL_ENV_INDICES, self._ALL_BODY_INDICES, device=self.device - ), - is_global=False, - ) + composer = self._permanent_wrench_composer + composer.compose_to_body_frame() + self.root_view.apply_forces_and_torques_at_position( + force_data=self.reshape_data_to_view_2d(composer.out_force_b, device=self.device).view(wp.float32), + torque_data=self.reshape_data_to_view_2d(composer.out_torque_b, device=self.device).view(wp.float32), + position_data=None, + indices=self._env_body_ids_to_view_ids( + self._ALL_ENV_INDICES, self._ALL_BODY_INDICES, device=self.device + ), + is_global=False, + ) self._instantaneous_wrench_composer.reset() def update(self, dt: float) -> None: diff --git a/source/isaaclab_physx/isaaclab_physx/renderers/isaac_rtx_renderer_utils.py b/source/isaaclab_physx/isaaclab_physx/renderers/isaac_rtx_renderer_utils.py index 032bf001c79f..719bf70890b0 100644 --- a/source/isaaclab_physx/isaaclab_physx/renderers/isaac_rtx_renderer_utils.py +++ b/source/isaaclab_physx/isaaclab_physx/renderers/isaac_rtx_renderer_utils.py @@ -9,16 +9,17 @@ import logging import time +from typing import Any import isaaclab.sim as sim_utils logger = logging.getLogger(__name__) -# Module-level dedup stamp: tracks the last (sim instance, physics step) at +# Module-level dedup stamp: tracks the last (sim instance, physics step, render generation) at # which Kit's ``app.update()`` was pumped. Keyed on ``id(sim)`` so that a # new ``SimulationContext`` (e.g. in a new test) automatically invalidates # any stale stamp from a previous instance. -_last_render_update_key: tuple[int, int] = (0, -1) +_last_render_update_key: tuple[int, int, int] = (0, -1, -1) _STREAMING_WAIT_TIMEOUT_S: float = 30.0 @@ -58,7 +59,7 @@ def _wait_for_streaming_complete() -> None: def ensure_isaac_rtx_render_update() -> None: - """Ensure the Isaac RTX renderer has been pumped for the current physics step. + """Ensure the Isaac RTX renderer has been pumped for the current sim step. This keeps the Kit-specific ``app.update()`` logic inside the renderers package rather than in the backend-agnostic ``SimulationContext``. @@ -66,11 +67,11 @@ def ensure_isaac_rtx_render_update() -> None: Safe to call from multiple ``Camera`` / ``TiledCamera`` instances per step — only the first call triggers ``app.update()``. Subsequent calls are no-ops because the module-level ``_last_render_update_key`` already matches the - current ``(id(sim), step_count)`` pair. + current ``(id(sim), step_count, render_generation)`` tuple. - The key is a ``(sim_instance_id, step_count)`` tuple so that creating a new - ``SimulationContext`` (e.g. in a subsequent test) automatically invalidates - any stale stamp left over from a previous instance. + The key is a ``(sim_instance_id, step_count, render_generation)`` tuple so that: + - creating a new ``SimulationContext`` invalidates stale stamps, and + - render/reset transitions that do not advance physics step count still force a fresh update. After the initial ``app.update()`` the streaming subsystem is queried synchronously via ``UsdContext.get_stage_streaming_status()``. If textures @@ -88,7 +89,8 @@ def ensure_isaac_rtx_render_update() -> None: if sim is None: return - key = (id(sim), sim._physics_step_count) + render_generation = getattr(sim, "render_generation", getattr(sim, "_render_generation", 0)) + key = (id(sim), sim._physics_step_count, render_generation) if _last_render_update_key == key: return # Already pumped this step (by another camera or a visualizer) @@ -116,3 +118,28 @@ def ensure_isaac_rtx_render_update() -> None: sim.set_setting("/app/player/playSimulations", True) _last_render_update_key = key + + +def pump_kit_app_for_headless_video_render_if_needed(sim: Any) -> None: + """Pump Kit app-loop for headless rgb-array rendering when needed. + + Isaac Sim / RTX specific; kept out of backend-agnostic :class:`~isaaclab.sim.SimulationContext`. + """ + if not bool(sim.get_setting("/isaaclab/video/enabled")): + return + + from isaaclab.utils.version import has_kit + + if not has_kit(): + return + if any(viz.pumps_app_update() for viz in sim.visualizers): + return + try: + ensure_isaac_rtx_render_update() + except (ImportError, AttributeError, ModuleNotFoundError) as exc: + logger.debug("[isaac_rtx] Skipping Kit app-loop pump in render() (non-Kit env): %s", exc) + except Exception as exc: + logger.warning( + "[isaac_rtx] Kit app-loop pump failed in render() — video frames may be stale or black: %s", + exc, + ) diff --git a/source/isaaclab_physx/isaaclab_physx/renderers/kit_viewport_utils.py b/source/isaaclab_physx/isaaclab_physx/renderers/kit_viewport_utils.py new file mode 100644 index 000000000000..af421a032399 --- /dev/null +++ b/source/isaaclab_physx/isaaclab_physx/renderers/kit_viewport_utils.py @@ -0,0 +1,35 @@ +# 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 + +"""Kit / Omniverse viewport helpers (Isaac Sim specific). + +These live in :mod:`isaaclab_physx` so :class:`~isaaclab.sim.SimulationContext` stays +backend-agnostic. +""" + +from __future__ import annotations + +import logging + +logger = logging.getLogger(__name__) + + +def set_kit_renderer_camera_view( + eye: tuple[float, float, float] | list[float], + target: tuple[float, float, float] | list[float], + camera_prim_path: str = "/OmniverseKit_Persp", +) -> None: + """Set camera view for the renderer/viewport camera only. + + This does not broadcast to visualizers. + """ + try: + import isaacsim.core.utils.viewports as isaacsim_viewports + + isaacsim_viewports.set_camera_view(eye=list(eye), target=list(target), camera_prim_path=str(camera_prim_path)) + except (ImportError, ModuleNotFoundError) as exc: + logger.debug("[kit_viewport] Renderer camera update skipped (no Kit): %s", exc) + except Exception as exc: + logger.warning("[kit_viewport] Renderer camera update failed: %s", exc) diff --git a/source/isaaclab_physx/isaaclab_physx/scene_data_providers/__init__.pyi b/source/isaaclab_physx/isaaclab_physx/scene_data_providers/__init__.pyi index d1612d1f3bbd..32c6f9c07335 100644 --- a/source/isaaclab_physx/isaaclab_physx/scene_data_providers/__init__.pyi +++ b/source/isaaclab_physx/isaaclab_physx/scene_data_providers/__init__.pyi @@ -4,9 +4,7 @@ # SPDX-License-Identifier: BSD-3-Clause __all__ = [ - "NewtonSceneDataProvider", "PhysxSceneDataProvider", ] -from .newton_scene_data_provider import NewtonSceneDataProvider from .physx_scene_data_provider import PhysxSceneDataProvider 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 9403cd40aa46..814d70a245bc 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 @@ -22,9 +22,6 @@ logger = logging.getLogger(__name__) -# Path pattern for env prims: /World/envs/env_/... -_ENV_ID_RE = re.compile(r"/World/envs/env_(\d+)") - @wp.kernel(enable_backward=False) def _set_body_q_kernel( @@ -37,36 +34,18 @@ def _set_body_q_kernel( body_q[i] = wp.transformf(positions[i], orientations[i]) -@wp.kernel(enable_backward=False) -def _set_body_q_subset_kernel( - positions: wp.array(dtype=wp.vec3), - orientations: wp.array(dtype=wp.quatf), - body_indices: wp.array(dtype=wp.int32), - body_q: wp.array(dtype=wp.transformf), -): - """Write pose arrays into selected Newton ``body_q`` indices.""" - i = wp.tid() - bi = body_indices[i] - body_q[bi] = wp.transformf(positions[i], orientations[i]) - - 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 + - Newton model/state (from the simulation context prebuilt payload when required) """ # ---- Environment discovery / metadata ------------------------------------------------- - def _env_id_from_path(self, path: str) -> int | None: - """Extract env id from path (e.g. /World/envs/env_42/...). Used to map body paths to envs for sync.""" - m = _ENV_ID_RE.search(path) - return int(m.group(1)) if m else None - def get_num_envs(self) -> int: """Return env count from stage discovery, cached once available.""" if self._num_envs is not None and self._num_envs > 0: @@ -118,10 +97,6 @@ def __init__(self, stage, simulation_context) -> None: requirements = self._simulation_context.get_scene_data_requirements() self._needs_newton_sync = bool(requirements.requires_newton_model) - # Benchmark/debug override: force USD traversal fallback even when prebuilt - # visualizer artifacts are available from the cloner path. - self._force_usd_fallback_for_newton_model_build = False - # Fixed metadata for visualizers. get_metadata() returns this plus num_envs so visualizers # can .get("num_envs", 0), .get("physics_backend", ...) etc. without the provider exposing many methods. self._metadata = {"physics_backend": "omni"} @@ -130,21 +105,14 @@ def __init__(self, stage, simulation_context) -> None: "[PhysxSceneDataProvider] USD stage is None and not available from simulation_context. " "Ensure the simulation context has a valid stage when using OV/Newton/Rerun/Viser visualizers." ) - self._up_axis = UsdGeom.GetStageUpAxis(self._stage) self._num_envs_at_last_newton_build: int | None = None # for _refresh_newton_model_if_needed self._device = getattr(self._simulation_context, "device", "cuda:0") self._newton_model = None self._newton_state = None - self._filtered_newton_model = None - self._filtered_newton_state = None - self._filtered_env_ids_key: tuple[int, ...] | None = None - self._filtered_body_indices: list[int] = [] self._rigid_body_paths: list[str] = [] # Paths used to create PhysX views. May include articulation roots for coverage. self._rigid_body_view_paths: list[str] = [] - # env_id -> list of body indices (in Newton body_key order) - self._env_id_to_body_indices: dict[int, list[int]] = {} # Reused pose buffers (MR perf): avoid per-call allocations in _read_poses_from_best_source. self._pose_buf_num_bodies = 0 @@ -154,14 +122,12 @@ def __init__(self, stage, simulation_context) -> None: self._xform_mask_buf = None # View index order as device tensors for vectorized scatter in _apply_view_poses. self._view_order_tensors: dict[str, Any] = {} - # Last full-model build source for tests/debugging ("prebuilt", "usd_fallback", "error"). + # Last load outcome (tests / debug): "prebuilt" | "missing" | "error". self._last_newton_model_build_source: str | None = None self._last_newton_model_build_elapsed_ms: float | None = None - # Initialize Newton pipeline only if needed for visualization if self._needs_newton_sync: - self._build_newton_model_from_usd() - self._build_env_id_to_body_indices() + self._load_newton_model_from_prebuilt_artifact() self._setup_rigid_body_view() # ---- Newton model + PhysX view setup -------------------------------------------------- @@ -174,7 +140,7 @@ def _wildcard_env_paths(self, paths: list[str]) -> list[str]: return list(dict.fromkeys(wildcard_paths)) if wildcard_paths else paths def _refresh_newton_model_if_needed(self) -> None: - """Rebuild Newton model/state and PhysX views if discovered env count changes.""" + """Reload Newton model/state and PhysX views when the discovered env count changes.""" num_envs = self.get_num_envs() if num_envs <= 0: return @@ -182,8 +148,7 @@ def _refresh_newton_model_if_needed(self) -> None: needs_rebuild = self._newton_model is None or self._newton_state is None needs_rebuild = needs_rebuild or (self._num_envs_at_last_newton_build != num_envs) if needs_rebuild: - self._build_newton_model_from_usd() - self._build_env_id_to_body_indices() + self._load_newton_model_from_prebuilt_artifact() self._setup_rigid_body_view() def _model_body_paths(self, model) -> list[str]: @@ -199,90 +164,45 @@ def _model_body_paths(self, model) -> list[str]: return [] return list(getattr(model, "body_label", None) or getattr(model, "body_key", [])) - def _try_use_prebuilt_newton_artifact(self) -> bool: - """Use scene-time prebuilt Newton visualizer artifact when available. - - Returns: - ``True`` when a valid prebuilt artifact was consumed, otherwise ``False``. - """ - if self._force_usd_fallback_for_newton_model_build: - return False - artifact = self._simulation_context.get_scene_data_visualizer_prebuilt_artifact() - if not artifact: - return False - - model = artifact.model - state = artifact.state - if model is None or state is None: - return False - - self._newton_model = model - self._newton_state = state - - # The Newton artifact was generated before all envs were cloned on the stage, so we update the shape colors - # in the Newton model here as the envs should have been cloned. - replace_newton_shape_colors(self._newton_model, self._stage) - - body_paths = list(artifact.rigid_body_paths) or self._model_body_paths(model) - # Keep one-to-one alignment between `body_paths` and Newton `state.body_q`. - # Articulation root prims are not body_q entries and must not be mixed here. - self._rigid_body_paths = body_paths - # Build the PhysX-view query set separately so articulation roots can still be - # included for view coverage without breaking body_q alignment. - view_paths = list(body_paths) - if artifact.articulation_paths: - seen = set(view_paths) - for path in artifact.articulation_paths: - if path not in seen: - view_paths.append(path) - seen.add(path) - self._rigid_body_view_paths = view_paths - self._xform_views.clear() - self._view_body_index_map = {} - self._view_order_tensors.clear() - self._pose_buf_num_bodies = 0 - self._positions_buf = None - self._orientations_buf = None - self._covered_buf = None - self._xform_mask_buf = None - self._env_id_to_body_indices = {} - self._num_envs_at_last_newton_build = int(artifact.num_envs) - self._filtered_newton_model = None - self._filtered_newton_state = None - self._filtered_env_ids_key = None - self._filtered_body_indices = [] - return True - - def _build_newton_model_from_usd(self) -> None: - """Build Newton model from USD and cache body paths.""" - # TODO: Deprecate this USD-traversal fallback once cloner/prebuilt coverage - # is complete for full and partial visualization model-build paths. + def _load_newton_model_from_prebuilt_artifact(self) -> None: + """Load Newton model and state from the simulation context prebuilt artifact.""" start_t = time.perf_counter() try: - if self._try_use_prebuilt_newton_artifact(): - self._last_newton_model_build_source = "prebuilt" + artifact = self._simulation_context.get_scene_data_visualizer_prebuilt_artifact() + if not artifact: + self._last_newton_model_build_source = "missing" + logger.error( + "[PhysxSceneDataProvider] No visualizer prebuilt artifact on the simulation context " + "(expected VisualizerPrebuiltArtifacts from scene setup)." + ) + self._clear_newton_model_state() return - self._last_newton_model_build_source = ( - "usd_fallback_forced" if self._force_usd_fallback_for_newton_model_build else "usd_fallback" - ) - from newton import ModelBuilder - builder = ModelBuilder(up_axis=self._up_axis) - builder.add_usd(self._stage, ignore_paths=[r"/World/envs/.*"]) - for env_id in range(self.get_num_envs()): - builder.begin_world() - builder.add_usd(self._stage, root_path=f"/World/envs/env_{env_id}") - builder.end_world() + model = artifact.model + state = artifact.state + if model is None or state is None: + self._last_newton_model_build_source = "missing" + logger.error( + "[PhysxSceneDataProvider] Prebuilt artifact is missing model or state; cannot sync PhysX to Newton." + ) + self._clear_newton_model_state() + return - self._newton_model = builder.finalize(device=self._device) - self._newton_state = self._newton_model.state() + self._newton_model = model + self._newton_state = state replace_newton_shape_colors(self._newton_model, self._stage) - # Extract scene structure from Newton model (single source of truth) - self._rigid_body_paths = self._model_body_paths(self._newton_model) - self._rigid_body_view_paths = list(self._rigid_body_paths) - + body_paths = list(artifact.rigid_body_paths) or self._model_body_paths(model) + self._rigid_body_paths = body_paths + view_paths = list(body_paths) + if artifact.articulation_paths: + seen = set(view_paths) + for path in artifact.articulation_paths: + if path not in seen: + view_paths.append(path) + seen.add(path) + self._rigid_body_view_paths = view_paths self._xform_views.clear() self._view_body_index_map = {} self._view_order_tensors.clear() @@ -291,28 +211,12 @@ def _build_newton_model_from_usd(self) -> None: self._orientations_buf = None self._covered_buf = None self._xform_mask_buf = None - self._env_id_to_body_indices = {} - self._num_envs_at_last_newton_build = self.get_num_envs() - # Invalidate any filtered model when full model changes. - self._filtered_newton_model = None - self._filtered_newton_state = None - self._filtered_env_ids_key = None - self._filtered_body_indices = [] - except ModuleNotFoundError as exc: - self._last_newton_model_build_source = "error" - logger.error( - "[PhysxSceneDataProvider] Newton module not available. " - "Install the Newton backend to use newton/rerun/viser visualizers." - ) - logger.debug(f"[PhysxSceneDataProvider] Newton import error: {exc}") + self._num_envs_at_last_newton_build = int(artifact.num_envs) + self._last_newton_model_build_source = "prebuilt" except Exception as exc: self._last_newton_model_build_source = "error" - logger.error(f"[PhysxSceneDataProvider] Failed to build Newton model from USD: {exc}") - self._newton_model = None - self._newton_state = None - self._rigid_body_paths = [] - self._rigid_body_view_paths = [] - self._num_envs_at_last_newton_build = None + logger.error("[PhysxSceneDataProvider] Failed to load Newton model from prebuilt artifact: %s", exc) + self._clear_newton_model_state() finally: elapsed_ms = (time.perf_counter() - start_t) * 1000.0 self._last_newton_model_build_elapsed_ms = elapsed_ms @@ -321,74 +225,19 @@ def _build_newton_model_from_usd(self) -> None: except Exception: num_envs = -1 logger.debug( - "[PhysxSceneDataProvider] Newton model build source=%s num_envs=%d elapsed_ms=%.2f", + "[PhysxSceneDataProvider] Newton model load source=%s num_envs=%d elapsed_ms=%.2f", self._last_newton_model_build_source, num_envs, elapsed_ms, ) - def _build_filtered_newton_model(self, env_ids: list[int]) -> None: - """Build Newton model/state for a subset of environments. - - Args: - env_ids: Environment ids to include in the subset model. - """ - # TODO: Deprecate this USD-traversal fallback once cloner/prebuilt coverage - # is complete for full and partial visualization model-build paths. - try: - from newton import ModelBuilder - - # Newton model building from USD with partial visualization does not currently use cloner, - # and falls back to slower USD-stage traversal. - builder = ModelBuilder(up_axis=self._up_axis) - builder.add_usd(self._stage, ignore_paths=[r"/World/envs/.*"]) - for env_id in env_ids: - builder.begin_world() - builder.add_usd(self._stage, root_path=f"/World/envs/env_{env_id}") - builder.end_world() - - self._filtered_newton_model = builder.finalize(device=self._device) - self._filtered_newton_state = self._filtered_newton_model.state() - - replace_newton_shape_colors(self._filtered_newton_model, self._stage) - - full_index_by_path = {path: i for i, path in enumerate(self._rigid_body_paths)} - filtered_paths = self._model_body_paths(self._filtered_newton_model) - self._filtered_body_indices = [] - missing = [] - for path in filtered_paths: - idx = full_index_by_path.get(path) - if idx is None: - missing.append(path) - else: - self._filtered_body_indices.append(idx) - if missing: - logger.warning( - "[PhysxSceneDataProvider] Filtered model contains %d bodies not in full model.", - len(missing), - ) - except ModuleNotFoundError as exc: - logger.error( - "[PhysxSceneDataProvider] Newton module not available. " - "Install the Newton backend to use newton/rerun/viser visualizers." - ) - logger.debug(f"[PhysxSceneDataProvider] Newton import error: {exc}") - self._filtered_newton_model = None - self._filtered_newton_state = None - self._filtered_body_indices = [] - except Exception as exc: - logger.error(f"[PhysxSceneDataProvider] Failed to build filtered Newton model from USD: {exc}") - self._filtered_newton_model = None - self._filtered_newton_state = None - self._filtered_body_indices = [] - - def _build_env_id_to_body_indices(self) -> None: - """Build mapping env_id -> list of body indices from rigid_body_paths.""" - self._env_id_to_body_indices = {} - for body_idx, path in enumerate(self._rigid_body_paths): - eid = self._env_id_from_path(path) - if eid is not None: - self._env_id_to_body_indices.setdefault(eid, []).append(body_idx) + def _clear_newton_model_state(self) -> None: + """Clear cached Newton model, state, and rigid-body path lists.""" + self._newton_model = None + self._newton_state = None + self._rigid_body_paths = [] + self._rigid_body_view_paths = [] + self._num_envs_at_last_newton_build = None def _setup_rigid_body_view(self) -> None: """Create PhysX RigidBodyView from Newton's body paths. @@ -547,7 +396,7 @@ def _apply_view_poses(self, view: Any, view_key: str, positions: Any, orientatio return newton_indices.numel() return 0 - # Fallback: Python loop when view does not fully cover or cache missing. + # Per-index path when the view does not fully cover bodies or the order cache is missing. count = 0 for newton_idx, view_idx in enumerate(order): if view_idx is not None and not covered[newton_idx]: @@ -559,13 +408,10 @@ 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). - - This is slower but more robust when PhysX views don't cover all bodies. - """ + """Fill remaining body poses using ``XformPrimView`` for prims not covered by the rigid-body view.""" import torch - from isaaclab.sim.views import XformPrimView + from isaaclab.sim.views import FrameView uncovered = torch.where(~covered)[0].cpu().tolist() if not uncovered: @@ -577,14 +423,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 @@ -595,7 +441,7 @@ def _apply_xform_poses(self, positions: Any, orientations: Any, covered: Any, xf if len(self._xform_view_failures) > 0: self._warn_once( "xform-fallback-failures", - "[PhysxSceneDataProvider] Xform fallback failed for %d body paths.", + "[PhysxSceneDataProvider] XformPrimView reads failed for %d body paths.", len(self._xform_view_failures), level=logging.DEBUG, ) @@ -604,7 +450,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. """ @@ -643,13 +489,15 @@ def _read_poses_from_best_source(self) -> tuple[Any, Any, str, Any] | None: covered = self._covered_buf xform_mask = self._xform_mask_buf - # Apply sources in preferred order: rigid bodies, then USD fallback. rigid_count = self._apply_view_poses(self._rigid_body_view, "rigid_body_view", positions, orientations, covered) xform_count = self._apply_xform_poses(positions, orientations, covered, xform_mask) if rigid_count == 0: self._warn_once( "rigid-source-unused", - "[PhysxSceneDataProvider] RigidBodyView did not provide any body transforms; using fallback sources.", + ( + "[PhysxSceneDataProvider] RigidBodyView returned no transforms; " + "filled from XformPrimView where needed." + ), level=logging.DEBUG, ) @@ -672,18 +520,10 @@ def _get_set_body_q_kernel(self): """Return module-level Warp kernel for writing transforms to Newton state.""" return _set_body_q_kernel - def _get_set_body_q_subset_kernel(self): - """Return module-level Warp kernel for subset writes.""" - return _set_body_q_subset_kernel - # ---- Newton state sync ---------------------------------------------------------------- - def update(self, env_ids: list[int] | None = None) -> None: - """Sync PhysX transforms to Newton state for visualization. - - When env_ids is not None, only body indices belonging to those envs are written - (partial sync). When None, all bodies are synced. - """ + def update(self) -> None: + """Sync PhysX transforms into the full Newton state (one kernel launch).""" if not self._needs_newton_sync or self._newton_state is None: return @@ -701,41 +541,15 @@ def update(self, env_ids: list[int] | None = None) -> None: positions_wp = wp.from_torch(positions.reshape(-1, 3), dtype=wp.vec3) orientations_wp = wp.from_torch(orientations_xyzw, dtype=wp.quatf) - if env_ids is None or not env_ids or not self._env_id_to_body_indices: - # Fast path: full state sync in one kernel launch. - set_body_q = self._get_set_body_q_kernel() - if set_body_q is None or positions_wp.shape[0] != self._newton_state.body_q.shape[0]: - return - wp.launch( - set_body_q, - dim=positions_wp.shape[0], - inputs=[positions_wp, orientations_wp, self._newton_state.body_q], - device=self._device, - ) - else: - body_indices = [] - for eid in env_ids: - body_indices.extend(self._env_id_to_body_indices.get(eid, [])) - if not body_indices: - return - # Subset path: write only env-selected body indices. - subset_kernel = self._get_set_body_q_subset_kernel() - if subset_kernel is None: - return - import torch - - indices_t = torch.tensor(body_indices, dtype=torch.int32, device=self._device) - pos_subset = positions.reshape(-1, 3)[body_indices] - ori_subset = orientations_xyzw[body_indices] - indices_wp = wp.from_torch(indices_t, dtype=wp.int32) - pos_wp = wp.from_torch(pos_subset.contiguous(), dtype=wp.vec3) - ori_wp = wp.from_torch(ori_subset.contiguous(), dtype=wp.quatf) - wp.launch( - subset_kernel, - dim=len(body_indices), - inputs=[pos_wp, ori_wp, indices_wp, self._newton_state.body_q], - device=self._device, - ) + set_body_q = self._get_set_body_q_kernel() + if set_body_q is None or positions_wp.shape[0] != self._newton_state.body_q.shape[0]: + return + wp.launch( + set_body_q, + dim=positions_wp.shape[0], + inputs=[positions_wp, orientations_wp, self._newton_state.body_q], + device=self._device, + ) except Exception as exc: self._warn_once( "newton-sync-update-failed", @@ -751,97 +565,11 @@ def get_newton_model(self) -> Any | None: """ return self._newton_model if self._needs_newton_sync else None - def get_newton_model_for_env_ids(self, env_ids: list[int] | None) -> Any | None: - """Return Newton model for selected environments. - - Args: - env_ids: Optional environment ids. ``None`` returns full model. - - Returns: - Full or filtered Newton model, or ``None`` when unavailable. - """ - if not self._needs_newton_sync: - return None - if env_ids is None: - return self._newton_model - env_ids_key = tuple(sorted(env_ids)) - if self._filtered_newton_model is None or self._filtered_env_ids_key != env_ids_key: - self._filtered_env_ids_key = env_ids_key - self._build_filtered_newton_model(list(env_ids_key)) - return self._filtered_newton_model - - def get_newton_state(self, env_ids: list[int] | None = None) -> Any | None: - """Return Newton state when sync is enabled. - - If env_ids is None, returns the full state. If env_ids is provided, returns a - state-like object whose body_q contains only the bodies for those envs (same order - as in the full model, for use with e.g. max_worlds=len(env_ids)). - """ + def get_newton_state(self) -> Any | None: + """Return full Newton state when sync is enabled.""" if not self._needs_newton_sync or self._newton_state is None: return None - if env_ids is None: - return self._newton_state - if not self._env_id_to_body_indices: - return self._create_empty_subset_state() - env_ids_key = tuple(sorted(env_ids)) - if self._filtered_newton_model is not None and self._filtered_env_ids_key == env_ids_key: - if not self._filtered_body_indices: - return self._create_empty_subset_state() - try: - import warp as wp - - body_q_t = wp.to_torch(self._newton_state.body_q) - subset = body_q_t[self._filtered_body_indices].clone() - self._filtered_newton_state.body_q = wp.from_torch(subset, dtype=wp.transformf) - return self._filtered_newton_state - except Exception: - return self._newton_state - body_indices = [] - for eid in env_ids: - body_indices.extend(self._env_id_to_body_indices.get(eid, [])) - if not body_indices: - return self._create_empty_subset_state() - - body_q = self._newton_state.body_q - try: - import warp as wp - - body_q_t = wp.to_torch(body_q) - body_q_subset = body_q_t[body_indices].clone() - except Exception: - return self._newton_state - return self._create_subset_state(body_q_subset) - - def _create_empty_subset_state(self): - """Return a minimal state-like object with empty body_q.""" - if self._newton_state is None: - return None - try: - import warp as wp - - body_q_t = wp.to_torch(self._newton_state.body_q) - empty = body_q_t[:0].clone() - return self._create_subset_state(empty) - except Exception: - return self._newton_state - - # ---- Newton subset helpers ------------------------------------------------------------- - - def _create_subset_state(self, body_q_subset): - """Return a minimal state-like object for subset rendering.""" - import warp as wp - - if hasattr(body_q_subset, "device") and not isinstance(body_q_subset, wp.array): - body_q_subset = wp.from_torch(body_q_subset, dtype=wp.transformf) - - class _SubsetState: - """Minimal state carrier with ``body_q`` field for subset rendering.""" - - pass - - s = _SubsetState() - s.body_q = body_q_subset - return s + return self._newton_state # ---- Public provider API --------------------------------------------------------------- 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/setup.py b/source/isaaclab_physx/setup.py index bc37a24d1694..1e917e938c2b 100644 --- a/source/isaaclab_physx/setup.py +++ b/source/isaaclab_physx/setup.py @@ -16,11 +16,13 @@ EXTENSION_TOML_DATA = toml.load(os.path.join(EXTENSION_PATH, "config", "extension.toml")) # Minimum dependencies required prior to installation -INSTALL_REQUIRES = [ - # INTENTIONALLY disabled to avoid circular dependency with isaaclab_physx, which also depends on isaaclab_newton. - # This will be re-enabled once we move to UV and pyproject.toml-based packaging. - # f"isaaclab_newton @ file://{os.path.join(os.path.dirname(EXTENSION_PATH), 'isaaclab_newton')}", -] +INSTALL_REQUIRES = [] + +EXTRAS_REQUIRE = { + "newton": [ + "newton @ git+https://github.com/newton-physics/newton.git@2684d75bfa4bb8b058a93b81c458a74b7701c997", + ], +} # Installation operation setup( @@ -36,6 +38,7 @@ package_data={"": ["*.pyi"]}, python_requires=">=3.12", install_requires=INSTALL_REQUIRES, + extras_require=EXTRAS_REQUIRE, packages=[ "isaaclab_physx", "isaaclab_physx.assets", 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 1a579ed0ef48..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.22" +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 3c9e06860e45..c81768405791 100644 --- a/source/isaaclab_tasks/docs/CHANGELOG.rst +++ b/source/isaaclab_tasks/docs/CHANGELOG.rst @@ -1,6 +1,30 @@ 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) +~~~~~~~~~~~~~~~~~~~ + +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/direct/anymal_c/anymal_c_env.py b/source/isaaclab_tasks/isaaclab_tasks/direct/anymal_c/anymal_c_env.py index ea95dfe5b98d..e58e377a0d63 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/direct/anymal_c/anymal_c_env.py +++ b/source/isaaclab_tasks/isaaclab_tasks/direct/anymal_c/anymal_c_env.py @@ -88,7 +88,9 @@ def _get_observations(self) -> dict: height_data = None if isinstance(self.cfg, AnymalCRoughEnvCfg): height_data = ( - self._height_scanner.data.pos_w[:, 2].unsqueeze(1) - self._height_scanner.data.ray_hits_w[..., 2] - 0.5 + wp.to_torch(self._height_scanner.data.pos_w)[:, 2].unsqueeze(1) + - wp.to_torch(self._height_scanner.data.ray_hits_w)[..., 2] + - 0.5 ).clip(-1.0, 1.0) obs = torch.cat( [ 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_tasks/isaaclab_tasks/utils/sim_launcher.py b/source/isaaclab_tasks/isaaclab_tasks/utils/sim_launcher.py index 964027858dae..0f0c6e5404de 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/utils/sim_launcher.py +++ b/source/isaaclab_tasks/isaaclab_tasks/utils/sim_launcher.py @@ -190,12 +190,38 @@ def launch_simulation( import importlib.util if importlib.util.find_spec("omni.kit") is None: + # Print a more obvious hint when a local _isaac_sim symlink + # exists but its env wasn't sourced (typical on Win11 + conda + # when activate.d hooks didn't fire, e.g. under `conda run`). + import os + import sys + + isaaclab_path = os.environ.get("ISAACLAB_PATH") + local_sim = os.path.join(isaaclab_path, "_isaac_sim") if isaaclab_path else None + extra_hint = "" + if local_sim and os.path.isdir(local_sim): + if sys.platform == "win32": + extra_hint = ( + f" Found a local Isaac Sim at {local_sim} but its environment is not active.\n" + f" Either run via `isaaclab.bat ...` (which now sources setup_conda_env.bat\n" + f" automatically), or in your current shell run:\n" + f' call "{local_sim}\\setup_conda_env.bat"\n' + ) + else: + extra_hint = ( + f" Found a local Isaac Sim at {local_sim} but its environment is not active.\n" + f" Either run via `./isaaclab.sh ...` (which now sources setup_conda_env.sh\n" + f" automatically), or in your current shell run:\n" + f' source "{local_sim}/setup_conda_env.sh"\n' + ) + logger.error( "\n[ERROR] Isaac Sim is not installed or not found on PYTHONPATH.\n" "\n" " This environment requires Isaac Sim and Omniverse Kit.\n" " PhysX backend and Kit visualizer currently requires Isaac Sim.\n" "\n" + f"{extra_hint}" " To fix this, ensure Isaac Sim is installed and available in the current environment.\n" "\n" " See https://isaac-sim.github.io/IsaacLab/main/source/setup/installation for details.\n" @@ -214,15 +240,17 @@ def launch_simulation( # Newton path without Kit: AppLauncher is skipped, so manually store the visualizer # selection in SettingsManager (works in standalone mode via plain dict) so that # SimulationContext._get_cli_visualizer_types() can find it. - from isaaclab.app.settings_manager import get_settings_manager + from isaaclab.app import AppLauncher disable_all = "none" in visualizer_types - active_types = [] if disable_all else sorted(visualizer_types) - visualizer_str = " ".join(active_types) - settings = get_settings_manager() - settings.set_string("/isaaclab/visualizer/types", visualizer_str) - settings.set_bool("/isaaclab/visualizer/explicit", True) - settings.set_bool("/isaaclab/visualizer/disable_all", disable_all) + if isinstance(launcher_args, argparse.Namespace): + AppLauncher.sync_visualizer_cli_settings_to_carb( + {**vars(launcher_args), "visualizer_explicit": True, "visualizer_disable_all": disable_all} + ) + elif isinstance(launcher_args, dict): + AppLauncher.sync_visualizer_cli_settings_to_carb( + {**launcher_args, "visualizer_explicit": True, "visualizer_disable_all": disable_all} + ) try: yield diff --git a/source/isaaclab_tasks/test/golden_images/cartpole/newton-newton_renderer-depth.png b/source/isaaclab_tasks/test/golden_images/cartpole/newton-newton_renderer-depth.png index 5853518fcd66..fd54026086e6 100644 Binary files a/source/isaaclab_tasks/test/golden_images/cartpole/newton-newton_renderer-depth.png and b/source/isaaclab_tasks/test/golden_images/cartpole/newton-newton_renderer-depth.png differ diff --git a/source/isaaclab_tasks/test/golden_images/cartpole/newton-newton_renderer-rgb.png b/source/isaaclab_tasks/test/golden_images/cartpole/newton-newton_renderer-rgb.png index 7835e4274cee..cd4ffb12e9a5 100644 Binary files a/source/isaaclab_tasks/test/golden_images/cartpole/newton-newton_renderer-rgb.png and b/source/isaaclab_tasks/test/golden_images/cartpole/newton-newton_renderer-rgb.png differ diff --git a/source/isaaclab_tasks/test/golden_images/cartpole/newton-newton_renderer-rgba.png b/source/isaaclab_tasks/test/golden_images/cartpole/newton-newton_renderer-rgba.png index 6e28bd204f3c..cf3e3dd4ddbc 100644 Binary files a/source/isaaclab_tasks/test/golden_images/cartpole/newton-newton_renderer-rgba.png and b/source/isaaclab_tasks/test/golden_images/cartpole/newton-newton_renderer-rgba.png differ diff --git a/source/isaaclab_tasks/test/golden_images/cartpole/physx-isaacsim_rtx_renderer-albedo.png b/source/isaaclab_tasks/test/golden_images/cartpole/physx-isaacsim_rtx_renderer-albedo.png index 28449eca2272..b43eaf1e120b 100644 Binary files a/source/isaaclab_tasks/test/golden_images/cartpole/physx-isaacsim_rtx_renderer-albedo.png and b/source/isaaclab_tasks/test/golden_images/cartpole/physx-isaacsim_rtx_renderer-albedo.png differ diff --git a/source/isaaclab_tasks/test/golden_images/cartpole/physx-isaacsim_rtx_renderer-depth.png b/source/isaaclab_tasks/test/golden_images/cartpole/physx-isaacsim_rtx_renderer-depth.png index 81e0489ca896..ab5b9cc096d0 100644 Binary files a/source/isaaclab_tasks/test/golden_images/cartpole/physx-isaacsim_rtx_renderer-depth.png and b/source/isaaclab_tasks/test/golden_images/cartpole/physx-isaacsim_rtx_renderer-depth.png differ diff --git a/source/isaaclab_tasks/test/golden_images/cartpole/physx-isaacsim_rtx_renderer-rgb.png b/source/isaaclab_tasks/test/golden_images/cartpole/physx-isaacsim_rtx_renderer-rgb.png index 9e3e93fe051a..6a5e22214631 100644 Binary files a/source/isaaclab_tasks/test/golden_images/cartpole/physx-isaacsim_rtx_renderer-rgb.png and b/source/isaaclab_tasks/test/golden_images/cartpole/physx-isaacsim_rtx_renderer-rgb.png differ diff --git a/source/isaaclab_tasks/test/golden_images/cartpole/physx-isaacsim_rtx_renderer-rgba.png b/source/isaaclab_tasks/test/golden_images/cartpole/physx-isaacsim_rtx_renderer-rgba.png index de26933c492f..c679cb05e9b3 100644 Binary files a/source/isaaclab_tasks/test/golden_images/cartpole/physx-isaacsim_rtx_renderer-rgba.png and b/source/isaaclab_tasks/test/golden_images/cartpole/physx-isaacsim_rtx_renderer-rgba.png differ diff --git a/source/isaaclab_tasks/test/golden_images/cartpole/physx-isaacsim_rtx_renderer-semantic_segmentation.png b/source/isaaclab_tasks/test/golden_images/cartpole/physx-isaacsim_rtx_renderer-semantic_segmentation.png index 013c6e0cbc95..a647622781cd 100644 Binary files a/source/isaaclab_tasks/test/golden_images/cartpole/physx-isaacsim_rtx_renderer-semantic_segmentation.png and b/source/isaaclab_tasks/test/golden_images/cartpole/physx-isaacsim_rtx_renderer-semantic_segmentation.png differ diff --git a/source/isaaclab_tasks/test/golden_images/cartpole/physx-isaacsim_rtx_renderer-simple_shading_constant_diffuse.png b/source/isaaclab_tasks/test/golden_images/cartpole/physx-isaacsim_rtx_renderer-simple_shading_constant_diffuse.png index 1c4916d96df5..effbdb581a9c 100644 Binary files a/source/isaaclab_tasks/test/golden_images/cartpole/physx-isaacsim_rtx_renderer-simple_shading_constant_diffuse.png and b/source/isaaclab_tasks/test/golden_images/cartpole/physx-isaacsim_rtx_renderer-simple_shading_constant_diffuse.png differ diff --git a/source/isaaclab_tasks/test/golden_images/cartpole/physx-isaacsim_rtx_renderer-simple_shading_diffuse_mdl.png b/source/isaaclab_tasks/test/golden_images/cartpole/physx-isaacsim_rtx_renderer-simple_shading_diffuse_mdl.png index 1c4916d96df5..effbdb581a9c 100644 Binary files a/source/isaaclab_tasks/test/golden_images/cartpole/physx-isaacsim_rtx_renderer-simple_shading_diffuse_mdl.png and b/source/isaaclab_tasks/test/golden_images/cartpole/physx-isaacsim_rtx_renderer-simple_shading_diffuse_mdl.png differ diff --git a/source/isaaclab_tasks/test/golden_images/cartpole/physx-isaacsim_rtx_renderer-simple_shading_full_mdl.png b/source/isaaclab_tasks/test/golden_images/cartpole/physx-isaacsim_rtx_renderer-simple_shading_full_mdl.png index 1c4916d96df5..effbdb581a9c 100644 Binary files a/source/isaaclab_tasks/test/golden_images/cartpole/physx-isaacsim_rtx_renderer-simple_shading_full_mdl.png and b/source/isaaclab_tasks/test/golden_images/cartpole/physx-isaacsim_rtx_renderer-simple_shading_full_mdl.png differ diff --git a/source/isaaclab_tasks/test/golden_images/cartpole/physx-newton_renderer-depth.png b/source/isaaclab_tasks/test/golden_images/cartpole/physx-newton_renderer-depth.png index 5853518fcd66..fd54026086e6 100644 Binary files a/source/isaaclab_tasks/test/golden_images/cartpole/physx-newton_renderer-depth.png and b/source/isaaclab_tasks/test/golden_images/cartpole/physx-newton_renderer-depth.png differ diff --git a/source/isaaclab_tasks/test/golden_images/cartpole/physx-newton_renderer-rgb.png b/source/isaaclab_tasks/test/golden_images/cartpole/physx-newton_renderer-rgb.png index 7835e4274cee..cd4ffb12e9a5 100644 Binary files a/source/isaaclab_tasks/test/golden_images/cartpole/physx-newton_renderer-rgb.png and b/source/isaaclab_tasks/test/golden_images/cartpole/physx-newton_renderer-rgb.png differ diff --git a/source/isaaclab_tasks/test/golden_images/cartpole/physx-newton_renderer-rgba.png b/source/isaaclab_tasks/test/golden_images/cartpole/physx-newton_renderer-rgba.png index 6e28bd204f3c..cf3e3dd4ddbc 100644 Binary files a/source/isaaclab_tasks/test/golden_images/cartpole/physx-newton_renderer-rgba.png and b/source/isaaclab_tasks/test/golden_images/cartpole/physx-newton_renderer-rgba.png differ diff --git a/source/isaaclab_tasks/test/golden_images/dexsuite_kuka/newton-isaacsim_rtx_renderer-albedo.png b/source/isaaclab_tasks/test/golden_images/dexsuite_kuka/newton-isaacsim_rtx_renderer-albedo.png old mode 100755 new mode 100644 index 45b171e2756b..458c3bda17d0 Binary files a/source/isaaclab_tasks/test/golden_images/dexsuite_kuka/newton-isaacsim_rtx_renderer-albedo.png and b/source/isaaclab_tasks/test/golden_images/dexsuite_kuka/newton-isaacsim_rtx_renderer-albedo.png differ diff --git a/source/isaaclab_tasks/test/golden_images/dexsuite_kuka/newton-isaacsim_rtx_renderer-depth.png b/source/isaaclab_tasks/test/golden_images/dexsuite_kuka/newton-isaacsim_rtx_renderer-depth.png old mode 100755 new mode 100644 diff --git a/source/isaaclab_tasks/test/golden_images/dexsuite_kuka/newton-isaacsim_rtx_renderer-rgb.png b/source/isaaclab_tasks/test/golden_images/dexsuite_kuka/newton-isaacsim_rtx_renderer-rgb.png index aedccbb16056..239700a5d10d 100644 Binary files a/source/isaaclab_tasks/test/golden_images/dexsuite_kuka/newton-isaacsim_rtx_renderer-rgb.png and b/source/isaaclab_tasks/test/golden_images/dexsuite_kuka/newton-isaacsim_rtx_renderer-rgb.png differ diff --git a/source/isaaclab_tasks/test/golden_images/dexsuite_kuka/newton-isaacsim_rtx_renderer-rgba.png b/source/isaaclab_tasks/test/golden_images/dexsuite_kuka/newton-isaacsim_rtx_renderer-rgba.png index 104e15d3abf2..d609ef2f5272 100644 Binary files a/source/isaaclab_tasks/test/golden_images/dexsuite_kuka/newton-isaacsim_rtx_renderer-rgba.png and b/source/isaaclab_tasks/test/golden_images/dexsuite_kuka/newton-isaacsim_rtx_renderer-rgba.png differ diff --git a/source/isaaclab_tasks/test/golden_images/dexsuite_kuka/newton-isaacsim_rtx_renderer-simple_shading_constant_diffuse.png b/source/isaaclab_tasks/test/golden_images/dexsuite_kuka/newton-isaacsim_rtx_renderer-simple_shading_constant_diffuse.png old mode 100755 new mode 100644 index 6c21c2da6195..c70ae2bcd604 Binary files a/source/isaaclab_tasks/test/golden_images/dexsuite_kuka/newton-isaacsim_rtx_renderer-simple_shading_constant_diffuse.png and b/source/isaaclab_tasks/test/golden_images/dexsuite_kuka/newton-isaacsim_rtx_renderer-simple_shading_constant_diffuse.png differ diff --git a/source/isaaclab_tasks/test/golden_images/dexsuite_kuka/newton-isaacsim_rtx_renderer-simple_shading_diffuse_mdl.png b/source/isaaclab_tasks/test/golden_images/dexsuite_kuka/newton-isaacsim_rtx_renderer-simple_shading_diffuse_mdl.png old mode 100755 new mode 100644 index cf24e0114464..b63ca3cf01d9 Binary files a/source/isaaclab_tasks/test/golden_images/dexsuite_kuka/newton-isaacsim_rtx_renderer-simple_shading_diffuse_mdl.png and b/source/isaaclab_tasks/test/golden_images/dexsuite_kuka/newton-isaacsim_rtx_renderer-simple_shading_diffuse_mdl.png differ diff --git a/source/isaaclab_tasks/test/golden_images/dexsuite_kuka/newton-isaacsim_rtx_renderer-simple_shading_full_mdl.png b/source/isaaclab_tasks/test/golden_images/dexsuite_kuka/newton-isaacsim_rtx_renderer-simple_shading_full_mdl.png old mode 100755 new mode 100644 index cf24e0114464..b63ca3cf01d9 Binary files a/source/isaaclab_tasks/test/golden_images/dexsuite_kuka/newton-isaacsim_rtx_renderer-simple_shading_full_mdl.png and b/source/isaaclab_tasks/test/golden_images/dexsuite_kuka/newton-isaacsim_rtx_renderer-simple_shading_full_mdl.png differ diff --git a/source/isaaclab_tasks/test/golden_images/dexsuite_kuka/newton-newton_renderer-depth.png b/source/isaaclab_tasks/test/golden_images/dexsuite_kuka/newton-newton_renderer-depth.png index dd932819cb01..cb4c89384aec 100644 Binary files a/source/isaaclab_tasks/test/golden_images/dexsuite_kuka/newton-newton_renderer-depth.png and b/source/isaaclab_tasks/test/golden_images/dexsuite_kuka/newton-newton_renderer-depth.png differ diff --git a/source/isaaclab_tasks/test/golden_images/dexsuite_kuka/newton-newton_renderer-rgb.png b/source/isaaclab_tasks/test/golden_images/dexsuite_kuka/newton-newton_renderer-rgb.png index 84e0672a0c92..034dcb4be991 100644 Binary files a/source/isaaclab_tasks/test/golden_images/dexsuite_kuka/newton-newton_renderer-rgb.png and b/source/isaaclab_tasks/test/golden_images/dexsuite_kuka/newton-newton_renderer-rgb.png differ diff --git a/source/isaaclab_tasks/test/golden_images/dexsuite_kuka/newton-newton_renderer-rgba.png b/source/isaaclab_tasks/test/golden_images/dexsuite_kuka/newton-newton_renderer-rgba.png index b0ea6a5629d5..6a8f543d4577 100644 Binary files a/source/isaaclab_tasks/test/golden_images/dexsuite_kuka/newton-newton_renderer-rgba.png and b/source/isaaclab_tasks/test/golden_images/dexsuite_kuka/newton-newton_renderer-rgba.png differ diff --git a/source/isaaclab_tasks/test/golden_images/dexsuite_kuka/physx-isaacsim_rtx_renderer-albedo.png b/source/isaaclab_tasks/test/golden_images/dexsuite_kuka/physx-isaacsim_rtx_renderer-albedo.png old mode 100755 new mode 100644 index e02f546f8ed2..7af151ef02c2 Binary files a/source/isaaclab_tasks/test/golden_images/dexsuite_kuka/physx-isaacsim_rtx_renderer-albedo.png and b/source/isaaclab_tasks/test/golden_images/dexsuite_kuka/physx-isaacsim_rtx_renderer-albedo.png differ diff --git a/source/isaaclab_tasks/test/golden_images/dexsuite_kuka/physx-isaacsim_rtx_renderer-depth.png b/source/isaaclab_tasks/test/golden_images/dexsuite_kuka/physx-isaacsim_rtx_renderer-depth.png old mode 100755 new mode 100644 diff --git a/source/isaaclab_tasks/test/golden_images/dexsuite_kuka/physx-isaacsim_rtx_renderer-rgb.png b/source/isaaclab_tasks/test/golden_images/dexsuite_kuka/physx-isaacsim_rtx_renderer-rgb.png index d1ea4fe5e781..de0d73165118 100644 Binary files a/source/isaaclab_tasks/test/golden_images/dexsuite_kuka/physx-isaacsim_rtx_renderer-rgb.png and b/source/isaaclab_tasks/test/golden_images/dexsuite_kuka/physx-isaacsim_rtx_renderer-rgb.png differ diff --git a/source/isaaclab_tasks/test/golden_images/dexsuite_kuka/physx-isaacsim_rtx_renderer-rgba.png b/source/isaaclab_tasks/test/golden_images/dexsuite_kuka/physx-isaacsim_rtx_renderer-rgba.png old mode 100755 new mode 100644 index dfe59520b408..a9e995427100 Binary files a/source/isaaclab_tasks/test/golden_images/dexsuite_kuka/physx-isaacsim_rtx_renderer-rgba.png and b/source/isaaclab_tasks/test/golden_images/dexsuite_kuka/physx-isaacsim_rtx_renderer-rgba.png differ diff --git a/source/isaaclab_tasks/test/golden_images/dexsuite_kuka/physx-isaacsim_rtx_renderer-semantic_segmentation.png b/source/isaaclab_tasks/test/golden_images/dexsuite_kuka/physx-isaacsim_rtx_renderer-semantic_segmentation.png old mode 100755 new mode 100644 diff --git a/source/isaaclab_tasks/test/golden_images/dexsuite_kuka/physx-isaacsim_rtx_renderer-simple_shading_constant_diffuse.png b/source/isaaclab_tasks/test/golden_images/dexsuite_kuka/physx-isaacsim_rtx_renderer-simple_shading_constant_diffuse.png old mode 100755 new mode 100644 index 784b4b983f7d..d2dba4b30b07 Binary files a/source/isaaclab_tasks/test/golden_images/dexsuite_kuka/physx-isaacsim_rtx_renderer-simple_shading_constant_diffuse.png and b/source/isaaclab_tasks/test/golden_images/dexsuite_kuka/physx-isaacsim_rtx_renderer-simple_shading_constant_diffuse.png differ diff --git a/source/isaaclab_tasks/test/golden_images/dexsuite_kuka/physx-isaacsim_rtx_renderer-simple_shading_diffuse_mdl.png b/source/isaaclab_tasks/test/golden_images/dexsuite_kuka/physx-isaacsim_rtx_renderer-simple_shading_diffuse_mdl.png old mode 100755 new mode 100644 index 784b4b983f7d..d2dba4b30b07 Binary files a/source/isaaclab_tasks/test/golden_images/dexsuite_kuka/physx-isaacsim_rtx_renderer-simple_shading_diffuse_mdl.png and b/source/isaaclab_tasks/test/golden_images/dexsuite_kuka/physx-isaacsim_rtx_renderer-simple_shading_diffuse_mdl.png differ diff --git a/source/isaaclab_tasks/test/golden_images/dexsuite_kuka/physx-isaacsim_rtx_renderer-simple_shading_full_mdl.png b/source/isaaclab_tasks/test/golden_images/dexsuite_kuka/physx-isaacsim_rtx_renderer-simple_shading_full_mdl.png old mode 100755 new mode 100644 index 417534e7f1c4..455f0ae674d5 Binary files a/source/isaaclab_tasks/test/golden_images/dexsuite_kuka/physx-isaacsim_rtx_renderer-simple_shading_full_mdl.png and b/source/isaaclab_tasks/test/golden_images/dexsuite_kuka/physx-isaacsim_rtx_renderer-simple_shading_full_mdl.png differ diff --git a/source/isaaclab_tasks/test/golden_images/dexsuite_kuka/physx-newton_renderer-depth.png b/source/isaaclab_tasks/test/golden_images/dexsuite_kuka/physx-newton_renderer-depth.png index f1cd28ce2e9e..044fa4e95faa 100644 Binary files a/source/isaaclab_tasks/test/golden_images/dexsuite_kuka/physx-newton_renderer-depth.png and b/source/isaaclab_tasks/test/golden_images/dexsuite_kuka/physx-newton_renderer-depth.png differ diff --git a/source/isaaclab_tasks/test/golden_images/dexsuite_kuka/physx-newton_renderer-rgb.png b/source/isaaclab_tasks/test/golden_images/dexsuite_kuka/physx-newton_renderer-rgb.png index 67cfc91228ab..2f17d09126c5 100644 Binary files a/source/isaaclab_tasks/test/golden_images/dexsuite_kuka/physx-newton_renderer-rgb.png and b/source/isaaclab_tasks/test/golden_images/dexsuite_kuka/physx-newton_renderer-rgb.png differ diff --git a/source/isaaclab_tasks/test/golden_images/dexsuite_kuka/physx-newton_renderer-rgba.png b/source/isaaclab_tasks/test/golden_images/dexsuite_kuka/physx-newton_renderer-rgba.png index 2eb1d3622f06..e59a3288aa64 100644 Binary files a/source/isaaclab_tasks/test/golden_images/dexsuite_kuka/physx-newton_renderer-rgba.png and b/source/isaaclab_tasks/test/golden_images/dexsuite_kuka/physx-newton_renderer-rgba.png differ diff --git a/source/isaaclab_tasks/test/golden_images/registered_tasks/Isaac-Cartpole-Albedo-Camera-Direct-v0/default_physics-default_renderer-albedo.png b/source/isaaclab_tasks/test/golden_images/registered_tasks/Isaac-Cartpole-Albedo-Camera-Direct-v0/default_physics-default_renderer-albedo.png index 84db641e9ee4..cfc89aa7ae2f 100644 Binary files a/source/isaaclab_tasks/test/golden_images/registered_tasks/Isaac-Cartpole-Albedo-Camera-Direct-v0/default_physics-default_renderer-albedo.png and b/source/isaaclab_tasks/test/golden_images/registered_tasks/Isaac-Cartpole-Albedo-Camera-Direct-v0/default_physics-default_renderer-albedo.png differ diff --git a/source/isaaclab_tasks/test/golden_images/registered_tasks/Isaac-Cartpole-Camera-Presets-Direct-v0/default_physics-default_renderer-rgb.png b/source/isaaclab_tasks/test/golden_images/registered_tasks/Isaac-Cartpole-Camera-Presets-Direct-v0/default_physics-default_renderer-rgb.png index 0a83503793d5..5e2f31e58d0d 100644 Binary files a/source/isaaclab_tasks/test/golden_images/registered_tasks/Isaac-Cartpole-Camera-Presets-Direct-v0/default_physics-default_renderer-rgb.png and b/source/isaaclab_tasks/test/golden_images/registered_tasks/Isaac-Cartpole-Camera-Presets-Direct-v0/default_physics-default_renderer-rgb.png differ diff --git a/source/isaaclab_tasks/test/golden_images/registered_tasks/Isaac-Cartpole-Camera-Presets-Direct-v0/default_physics-default_renderer-rgba.png b/source/isaaclab_tasks/test/golden_images/registered_tasks/Isaac-Cartpole-Camera-Presets-Direct-v0/default_physics-default_renderer-rgba.png index 4be4bad97e72..25bdb10e17a8 100644 Binary files a/source/isaaclab_tasks/test/golden_images/registered_tasks/Isaac-Cartpole-Camera-Presets-Direct-v0/default_physics-default_renderer-rgba.png and b/source/isaaclab_tasks/test/golden_images/registered_tasks/Isaac-Cartpole-Camera-Presets-Direct-v0/default_physics-default_renderer-rgba.png differ diff --git a/source/isaaclab_tasks/test/golden_images/registered_tasks/Isaac-Cartpole-Depth-Camera-Direct-v0/default_physics-default_renderer-depth.png b/source/isaaclab_tasks/test/golden_images/registered_tasks/Isaac-Cartpole-Depth-Camera-Direct-v0/default_physics-default_renderer-depth.png index 81e0489ca896..ab5b9cc096d0 100644 Binary files a/source/isaaclab_tasks/test/golden_images/registered_tasks/Isaac-Cartpole-Depth-Camera-Direct-v0/default_physics-default_renderer-depth.png and b/source/isaaclab_tasks/test/golden_images/registered_tasks/Isaac-Cartpole-Depth-Camera-Direct-v0/default_physics-default_renderer-depth.png differ diff --git a/source/isaaclab_tasks/test/golden_images/registered_tasks/Isaac-Cartpole-RGB-Camera-Direct-v0/default_physics-default_renderer-rgb.png b/source/isaaclab_tasks/test/golden_images/registered_tasks/Isaac-Cartpole-RGB-Camera-Direct-v0/default_physics-default_renderer-rgb.png index 2a9d492e5e57..5e2f31e58d0d 100644 Binary files a/source/isaaclab_tasks/test/golden_images/registered_tasks/Isaac-Cartpole-RGB-Camera-Direct-v0/default_physics-default_renderer-rgb.png and b/source/isaaclab_tasks/test/golden_images/registered_tasks/Isaac-Cartpole-RGB-Camera-Direct-v0/default_physics-default_renderer-rgb.png differ diff --git a/source/isaaclab_tasks/test/golden_images/registered_tasks/Isaac-Cartpole-RGB-Camera-Direct-v0/default_physics-default_renderer-rgba.png b/source/isaaclab_tasks/test/golden_images/registered_tasks/Isaac-Cartpole-RGB-Camera-Direct-v0/default_physics-default_renderer-rgba.png index e1bea2aadd79..25bdb10e17a8 100644 Binary files a/source/isaaclab_tasks/test/golden_images/registered_tasks/Isaac-Cartpole-RGB-Camera-Direct-v0/default_physics-default_renderer-rgba.png and b/source/isaaclab_tasks/test/golden_images/registered_tasks/Isaac-Cartpole-RGB-Camera-Direct-v0/default_physics-default_renderer-rgba.png differ diff --git a/source/isaaclab_tasks/test/golden_images/registered_tasks/Isaac-Cartpole-SimpleShading-Constant-Camera-Direct-v0/default_physics-default_renderer-simple_shading_constant_diffuse.png b/source/isaaclab_tasks/test/golden_images/registered_tasks/Isaac-Cartpole-SimpleShading-Constant-Camera-Direct-v0/default_physics-default_renderer-simple_shading_constant_diffuse.png index c9e291ad5912..d944fb4949af 100644 Binary files a/source/isaaclab_tasks/test/golden_images/registered_tasks/Isaac-Cartpole-SimpleShading-Constant-Camera-Direct-v0/default_physics-default_renderer-simple_shading_constant_diffuse.png and b/source/isaaclab_tasks/test/golden_images/registered_tasks/Isaac-Cartpole-SimpleShading-Constant-Camera-Direct-v0/default_physics-default_renderer-simple_shading_constant_diffuse.png differ diff --git a/source/isaaclab_tasks/test/golden_images/registered_tasks/Isaac-Cartpole-SimpleShading-Diffuse-Camera-Direct-v0/default_physics-default_renderer-simple_shading_diffuse_mdl.png b/source/isaaclab_tasks/test/golden_images/registered_tasks/Isaac-Cartpole-SimpleShading-Diffuse-Camera-Direct-v0/default_physics-default_renderer-simple_shading_diffuse_mdl.png index c9e291ad5912..d944fb4949af 100644 Binary files a/source/isaaclab_tasks/test/golden_images/registered_tasks/Isaac-Cartpole-SimpleShading-Diffuse-Camera-Direct-v0/default_physics-default_renderer-simple_shading_diffuse_mdl.png and b/source/isaaclab_tasks/test/golden_images/registered_tasks/Isaac-Cartpole-SimpleShading-Diffuse-Camera-Direct-v0/default_physics-default_renderer-simple_shading_diffuse_mdl.png differ diff --git a/source/isaaclab_tasks/test/golden_images/registered_tasks/Isaac-Cartpole-SimpleShading-Full-Camera-Direct-v0/default_physics-default_renderer-simple_shading_full_mdl.png b/source/isaaclab_tasks/test/golden_images/registered_tasks/Isaac-Cartpole-SimpleShading-Full-Camera-Direct-v0/default_physics-default_renderer-simple_shading_full_mdl.png index c9e291ad5912..d944fb4949af 100644 Binary files a/source/isaaclab_tasks/test/golden_images/registered_tasks/Isaac-Cartpole-SimpleShading-Full-Camera-Direct-v0/default_physics-default_renderer-simple_shading_full_mdl.png and b/source/isaaclab_tasks/test/golden_images/registered_tasks/Isaac-Cartpole-SimpleShading-Full-Camera-Direct-v0/default_physics-default_renderer-simple_shading_full_mdl.png differ diff --git a/source/isaaclab_tasks/test/golden_images/registered_tasks/Isaac-Repose-Cube-Shadow-Vision-Direct-v0/default_physics-default_renderer-depth.png b/source/isaaclab_tasks/test/golden_images/registered_tasks/Isaac-Repose-Cube-Shadow-Vision-Direct-v0/default_physics-default_renderer-depth.png index fb52a8e18211..b538c6dfdfa2 100644 Binary files a/source/isaaclab_tasks/test/golden_images/registered_tasks/Isaac-Repose-Cube-Shadow-Vision-Direct-v0/default_physics-default_renderer-depth.png and b/source/isaaclab_tasks/test/golden_images/registered_tasks/Isaac-Repose-Cube-Shadow-Vision-Direct-v0/default_physics-default_renderer-depth.png differ diff --git a/source/isaaclab_tasks/test/golden_images/registered_tasks/Isaac-Repose-Cube-Shadow-Vision-Direct-v0/default_physics-default_renderer-rgb.png b/source/isaaclab_tasks/test/golden_images/registered_tasks/Isaac-Repose-Cube-Shadow-Vision-Direct-v0/default_physics-default_renderer-rgb.png index 1936f5b32056..415da3d1a9c9 100644 Binary files a/source/isaaclab_tasks/test/golden_images/registered_tasks/Isaac-Repose-Cube-Shadow-Vision-Direct-v0/default_physics-default_renderer-rgb.png and b/source/isaaclab_tasks/test/golden_images/registered_tasks/Isaac-Repose-Cube-Shadow-Vision-Direct-v0/default_physics-default_renderer-rgb.png differ diff --git a/source/isaaclab_tasks/test/golden_images/registered_tasks/Isaac-Repose-Cube-Shadow-Vision-Direct-v0/default_physics-default_renderer-rgba.png b/source/isaaclab_tasks/test/golden_images/registered_tasks/Isaac-Repose-Cube-Shadow-Vision-Direct-v0/default_physics-default_renderer-rgba.png index 54057a527807..5e21ebb08121 100644 Binary files a/source/isaaclab_tasks/test/golden_images/registered_tasks/Isaac-Repose-Cube-Shadow-Vision-Direct-v0/default_physics-default_renderer-rgba.png and b/source/isaaclab_tasks/test/golden_images/registered_tasks/Isaac-Repose-Cube-Shadow-Vision-Direct-v0/default_physics-default_renderer-rgba.png differ diff --git a/source/isaaclab_tasks/test/golden_images/registered_tasks/Isaac-Repose-Cube-Shadow-Vision-Direct-v0/default_physics-default_renderer-semantic_segmentation.png b/source/isaaclab_tasks/test/golden_images/registered_tasks/Isaac-Repose-Cube-Shadow-Vision-Direct-v0/default_physics-default_renderer-semantic_segmentation.png index 2d34c534c973..3bca5a971ff0 100644 Binary files a/source/isaaclab_tasks/test/golden_images/registered_tasks/Isaac-Repose-Cube-Shadow-Vision-Direct-v0/default_physics-default_renderer-semantic_segmentation.png and b/source/isaaclab_tasks/test/golden_images/registered_tasks/Isaac-Repose-Cube-Shadow-Vision-Direct-v0/default_physics-default_renderer-semantic_segmentation.png differ diff --git a/source/isaaclab_tasks/test/golden_images/shadow_hand/newton-isaacsim_rtx_renderer-albedo.png b/source/isaaclab_tasks/test/golden_images/shadow_hand/newton-isaacsim_rtx_renderer-albedo.png old mode 100755 new mode 100644 diff --git a/source/isaaclab_tasks/test/golden_images/shadow_hand/newton-isaacsim_rtx_renderer-rgb.png b/source/isaaclab_tasks/test/golden_images/shadow_hand/newton-isaacsim_rtx_renderer-rgb.png index 756b7791329e..3072c8e24779 100644 Binary files a/source/isaaclab_tasks/test/golden_images/shadow_hand/newton-isaacsim_rtx_renderer-rgb.png and b/source/isaaclab_tasks/test/golden_images/shadow_hand/newton-isaacsim_rtx_renderer-rgb.png differ diff --git a/source/isaaclab_tasks/test/golden_images/shadow_hand/newton-isaacsim_rtx_renderer-rgba.png b/source/isaaclab_tasks/test/golden_images/shadow_hand/newton-isaacsim_rtx_renderer-rgba.png old mode 100755 new mode 100644 index 5575d93246cf..83632fa9d89a Binary files a/source/isaaclab_tasks/test/golden_images/shadow_hand/newton-isaacsim_rtx_renderer-rgba.png and b/source/isaaclab_tasks/test/golden_images/shadow_hand/newton-isaacsim_rtx_renderer-rgba.png differ diff --git a/source/isaaclab_tasks/test/golden_images/shadow_hand/newton-isaacsim_rtx_renderer-semantic_segmentation.png b/source/isaaclab_tasks/test/golden_images/shadow_hand/newton-isaacsim_rtx_renderer-semantic_segmentation.png index f3e2cdedae07..ffb770a4a555 100644 Binary files a/source/isaaclab_tasks/test/golden_images/shadow_hand/newton-isaacsim_rtx_renderer-semantic_segmentation.png and b/source/isaaclab_tasks/test/golden_images/shadow_hand/newton-isaacsim_rtx_renderer-semantic_segmentation.png differ diff --git a/source/isaaclab_tasks/test/golden_images/shadow_hand/newton-isaacsim_rtx_renderer-simple_shading_constant_diffuse.png b/source/isaaclab_tasks/test/golden_images/shadow_hand/newton-isaacsim_rtx_renderer-simple_shading_constant_diffuse.png old mode 100755 new mode 100644 diff --git a/source/isaaclab_tasks/test/golden_images/shadow_hand/newton-isaacsim_rtx_renderer-simple_shading_diffuse_mdl.png b/source/isaaclab_tasks/test/golden_images/shadow_hand/newton-isaacsim_rtx_renderer-simple_shading_diffuse_mdl.png old mode 100755 new mode 100644 diff --git a/source/isaaclab_tasks/test/golden_images/shadow_hand/newton-isaacsim_rtx_renderer-simple_shading_full_mdl.png b/source/isaaclab_tasks/test/golden_images/shadow_hand/newton-isaacsim_rtx_renderer-simple_shading_full_mdl.png old mode 100755 new mode 100644 index 7146c2359400..a4b844930dc2 Binary files a/source/isaaclab_tasks/test/golden_images/shadow_hand/newton-isaacsim_rtx_renderer-simple_shading_full_mdl.png and b/source/isaaclab_tasks/test/golden_images/shadow_hand/newton-isaacsim_rtx_renderer-simple_shading_full_mdl.png differ diff --git a/source/isaaclab_tasks/test/golden_images/shadow_hand/newton-newton_renderer-depth.png b/source/isaaclab_tasks/test/golden_images/shadow_hand/newton-newton_renderer-depth.png index cdc12e1ce228..720dbbf9f140 100644 Binary files a/source/isaaclab_tasks/test/golden_images/shadow_hand/newton-newton_renderer-depth.png and b/source/isaaclab_tasks/test/golden_images/shadow_hand/newton-newton_renderer-depth.png differ diff --git a/source/isaaclab_tasks/test/golden_images/shadow_hand/newton-newton_renderer-rgb.png b/source/isaaclab_tasks/test/golden_images/shadow_hand/newton-newton_renderer-rgb.png index 071e6d9b2dd2..b8ceef3ec3a5 100644 Binary files a/source/isaaclab_tasks/test/golden_images/shadow_hand/newton-newton_renderer-rgb.png and b/source/isaaclab_tasks/test/golden_images/shadow_hand/newton-newton_renderer-rgb.png differ diff --git a/source/isaaclab_tasks/test/golden_images/shadow_hand/newton-newton_renderer-rgba.png b/source/isaaclab_tasks/test/golden_images/shadow_hand/newton-newton_renderer-rgba.png index eb8542490adc..4e08fb2bbbc1 100644 Binary files a/source/isaaclab_tasks/test/golden_images/shadow_hand/newton-newton_renderer-rgba.png and b/source/isaaclab_tasks/test/golden_images/shadow_hand/newton-newton_renderer-rgba.png differ diff --git a/source/isaaclab_tasks/test/golden_images/shadow_hand/physx-isaacsim_rtx_renderer-albedo.png b/source/isaaclab_tasks/test/golden_images/shadow_hand/physx-isaacsim_rtx_renderer-albedo.png old mode 100755 new mode 100644 index c45e91b12616..ecf5eb5becf5 Binary files a/source/isaaclab_tasks/test/golden_images/shadow_hand/physx-isaacsim_rtx_renderer-albedo.png and b/source/isaaclab_tasks/test/golden_images/shadow_hand/physx-isaacsim_rtx_renderer-albedo.png differ diff --git a/source/isaaclab_tasks/test/golden_images/shadow_hand/physx-isaacsim_rtx_renderer-depth.png b/source/isaaclab_tasks/test/golden_images/shadow_hand/physx-isaacsim_rtx_renderer-depth.png index fb52a8e18211..b538c6dfdfa2 100644 Binary files a/source/isaaclab_tasks/test/golden_images/shadow_hand/physx-isaacsim_rtx_renderer-depth.png and b/source/isaaclab_tasks/test/golden_images/shadow_hand/physx-isaacsim_rtx_renderer-depth.png differ diff --git a/source/isaaclab_tasks/test/golden_images/shadow_hand/physx-isaacsim_rtx_renderer-rgb.png b/source/isaaclab_tasks/test/golden_images/shadow_hand/physx-isaacsim_rtx_renderer-rgb.png index 8803719c8afd..37c6a747c8c4 100644 Binary files a/source/isaaclab_tasks/test/golden_images/shadow_hand/physx-isaacsim_rtx_renderer-rgb.png and b/source/isaaclab_tasks/test/golden_images/shadow_hand/physx-isaacsim_rtx_renderer-rgb.png differ diff --git a/source/isaaclab_tasks/test/golden_images/shadow_hand/physx-isaacsim_rtx_renderer-rgba.png b/source/isaaclab_tasks/test/golden_images/shadow_hand/physx-isaacsim_rtx_renderer-rgba.png old mode 100755 new mode 100644 index 7b7a92e5832b..82ea7f71ced0 Binary files a/source/isaaclab_tasks/test/golden_images/shadow_hand/physx-isaacsim_rtx_renderer-rgba.png and b/source/isaaclab_tasks/test/golden_images/shadow_hand/physx-isaacsim_rtx_renderer-rgba.png differ diff --git a/source/isaaclab_tasks/test/golden_images/shadow_hand/physx-isaacsim_rtx_renderer-semantic_segmentation.png b/source/isaaclab_tasks/test/golden_images/shadow_hand/physx-isaacsim_rtx_renderer-semantic_segmentation.png old mode 100755 new mode 100644 index fc09a2e3eb62..3bca5a971ff0 Binary files a/source/isaaclab_tasks/test/golden_images/shadow_hand/physx-isaacsim_rtx_renderer-semantic_segmentation.png and b/source/isaaclab_tasks/test/golden_images/shadow_hand/physx-isaacsim_rtx_renderer-semantic_segmentation.png differ diff --git a/source/isaaclab_tasks/test/golden_images/shadow_hand/physx-isaacsim_rtx_renderer-simple_shading_constant_diffuse.png b/source/isaaclab_tasks/test/golden_images/shadow_hand/physx-isaacsim_rtx_renderer-simple_shading_constant_diffuse.png old mode 100755 new mode 100644 index d4f6aa7b4ce0..79819618b045 Binary files a/source/isaaclab_tasks/test/golden_images/shadow_hand/physx-isaacsim_rtx_renderer-simple_shading_constant_diffuse.png and b/source/isaaclab_tasks/test/golden_images/shadow_hand/physx-isaacsim_rtx_renderer-simple_shading_constant_diffuse.png differ diff --git a/source/isaaclab_tasks/test/golden_images/shadow_hand/physx-isaacsim_rtx_renderer-simple_shading_diffuse_mdl.png b/source/isaaclab_tasks/test/golden_images/shadow_hand/physx-isaacsim_rtx_renderer-simple_shading_diffuse_mdl.png old mode 100755 new mode 100644 index d4f6aa7b4ce0..79819618b045 Binary files a/source/isaaclab_tasks/test/golden_images/shadow_hand/physx-isaacsim_rtx_renderer-simple_shading_diffuse_mdl.png and b/source/isaaclab_tasks/test/golden_images/shadow_hand/physx-isaacsim_rtx_renderer-simple_shading_diffuse_mdl.png differ diff --git a/source/isaaclab_tasks/test/golden_images/shadow_hand/physx-isaacsim_rtx_renderer-simple_shading_full_mdl.png b/source/isaaclab_tasks/test/golden_images/shadow_hand/physx-isaacsim_rtx_renderer-simple_shading_full_mdl.png old mode 100755 new mode 100644 index d4f6aa7b4ce0..79819618b045 Binary files a/source/isaaclab_tasks/test/golden_images/shadow_hand/physx-isaacsim_rtx_renderer-simple_shading_full_mdl.png and b/source/isaaclab_tasks/test/golden_images/shadow_hand/physx-isaacsim_rtx_renderer-simple_shading_full_mdl.png differ diff --git a/source/isaaclab_tasks/test/golden_images/shadow_hand/physx-newton_renderer-depth.png b/source/isaaclab_tasks/test/golden_images/shadow_hand/physx-newton_renderer-depth.png index 99b3685ba904..8abb9f167e57 100644 Binary files a/source/isaaclab_tasks/test/golden_images/shadow_hand/physx-newton_renderer-depth.png and b/source/isaaclab_tasks/test/golden_images/shadow_hand/physx-newton_renderer-depth.png differ diff --git a/source/isaaclab_tasks/test/golden_images/shadow_hand/physx-newton_renderer-rgb.png b/source/isaaclab_tasks/test/golden_images/shadow_hand/physx-newton_renderer-rgb.png index 6a22c884e0b2..bacc0fcc47cb 100644 Binary files a/source/isaaclab_tasks/test/golden_images/shadow_hand/physx-newton_renderer-rgb.png and b/source/isaaclab_tasks/test/golden_images/shadow_hand/physx-newton_renderer-rgb.png differ diff --git a/source/isaaclab_tasks/test/golden_images/shadow_hand/physx-newton_renderer-rgba.png b/source/isaaclab_tasks/test/golden_images/shadow_hand/physx-newton_renderer-rgba.png index ba8907a13c3b..8772c6f51e42 100644 Binary files a/source/isaaclab_tasks/test/golden_images/shadow_hand/physx-newton_renderer-rgba.png and b/source/isaaclab_tasks/test/golden_images/shadow_hand/physx-newton_renderer-rgba.png differ diff --git a/source/isaaclab_tasks/test/test_rendering_correctness.py b/source/isaaclab_tasks/test/test_rendering_correctness.py index 3d3377901916..24ee5af0fac8 100644 --- a/source/isaaclab_tasks/test/test_rendering_correctness.py +++ b/source/isaaclab_tasks/test/test_rendering_correctness.py @@ -52,6 +52,15 @@ # _PIXEL_L2_NORM_DIFFERENCE_THRESHOLD = 10.0 +# The max percentage of pixels allowed to differ. If the percentage exceeds this value, the test will fail. +# The value is set case by case based on the screen space taken up by the env in camera output images. It +# needs to be large enough to tolerate minor rendering noise while small enough to catch unexpected changes. +_MAX_DIFFERENT_PIXELS_PERCENTAGE_BY_ENV_NAME = { + "cartpole": 1.0, + "shadow_hand": 3.0, + "dexsuite_kuka": 4.0, +} + _OVRTX_DISABLED = pytest.mark.skip( reason="OVRTX is optional and experimental feature and temporarily is excluded from testing." ) @@ -66,9 +75,6 @@ # "img_result_path": str | None, "img_golden_path": str | None} _COMPARISON_SCORES: list[dict] = [] -# Environment seed. -_ENV_SEED = 42 - # --------------------------------------------------------------------------- # Fixtures @@ -643,7 +649,6 @@ def shadow_hand_env(request): env_cfg = _apply_overrides_to_env_cfg(env_cfg, override_args) env_cfg.scene.num_envs = 4 - env_cfg.seed = _ENV_SEED if data_type == "depth": # Disable CNN forward pass as it cannot be meaningfully trained from depth alone and will raise a ValueError. @@ -652,7 +657,6 @@ def shadow_hand_env(request): env = None try: env = ShadowHandVisionEnv(env_cfg) - env.reset(seed=_ENV_SEED) yield physics_backend, renderer, data_type, env finally: if env is not None: @@ -662,13 +666,13 @@ def shadow_hand_env(request): def test_shadow_hand(shadow_hand_env): """Camera output must contain at least one non-zero pixel (Shadow Hand vision env).""" physics_backend, renderer, _, env = shadow_hand_env - + test_name = "shadow_hand" _validate_camera_outputs( - "shadow_hand", + test_name, physics_backend, renderer, env._tiled_camera.data.output, - max_different_pixels_percentage=8.0, + max_different_pixels_percentage=_MAX_DIFFERENT_PIXELS_PERCENTAGE_BY_ENV_NAME[test_name], ) @@ -691,12 +695,10 @@ def cartpole_env(request): env_cfg = _apply_overrides_to_env_cfg(env_cfg, override_args) env_cfg.scene.num_envs = 4 - env_cfg.seed = _ENV_SEED env = None try: env = CartpoleCameraEnv(env_cfg) - env.reset(seed=_ENV_SEED) yield physics_backend, renderer, data_type, env finally: if env is not None: @@ -706,13 +708,13 @@ def cartpole_env(request): def test_cartpole(cartpole_env): """Camera output must contain at least one non-zero pixel (Cartpole camera env).""" physics_backend, renderer, _, env = cartpole_env - + test_name = "cartpole" _validate_camera_outputs( - "cartpole", + test_name, physics_backend, renderer, env._tiled_camera.data.output, - max_different_pixels_percentage=2.0, + max_different_pixels_percentage=_MAX_DIFFERENT_PIXELS_PERCENTAGE_BY_ENV_NAME[test_name], ) @@ -739,12 +741,10 @@ def dexsuite_kuka_allegro_lift_env(request): env_cfg = _apply_overrides_to_env_cfg(env_cfg, override_args) env_cfg.scene.num_envs = 4 - env_cfg.seed = _ENV_SEED env = None try: env = ManagerBasedRLEnv(env_cfg) - env.reset(seed=_ENV_SEED) yield physics_backend, renderer, data_type, env finally: if env is not None: @@ -754,13 +754,13 @@ def dexsuite_kuka_allegro_lift_env(request): def test_dexsuite_kuka_allegro_lift(dexsuite_kuka_allegro_lift_env): """Camera output must contain at least one non-zero pixel (Dexsuite Kuka-Allegro Lift, single camera).""" physics_backend, renderer, _, env = dexsuite_kuka_allegro_lift_env - + test_name = "dexsuite_kuka" _validate_camera_outputs( - "dexsuite_kuka", + test_name, physics_backend, renderer, env.scene.sensors["base_camera"].data.output, - max_different_pixels_percentage=10.0, + max_different_pixels_percentage=_MAX_DIFFERENT_PIXELS_PERCENTAGE_BY_ENV_NAME[test_name], ) @@ -769,25 +769,25 @@ def test_dexsuite_kuka_allegro_lift(dexsuite_kuka_allegro_lift_env): # --------------------------------------------------------------------------- # Task IDs that expose camera/tiled_camera image observations; each is validated for non-blank rendering. +# The max different pixels percentage is set based on the screen space taken up by the env. _RENDER_CORRECTNESS_TASK_IDS = [ - "Isaac-Cartpole-Albedo-Camera-Direct-v0", - "Isaac-Cartpole-Camera-Presets-Direct-v0", - "Isaac-Cartpole-Depth-Camera-Direct-v0", - "Isaac-Cartpole-RGB-Camera-Direct-v0", - "Isaac-Cartpole-SimpleShading-Constant-Camera-Direct-v0", - "Isaac-Cartpole-SimpleShading-Diffuse-Camera-Direct-v0", - "Isaac-Cartpole-SimpleShading-Full-Camera-Direct-v0", - "Isaac-Repose-Cube-Shadow-Vision-Direct-v0", + ("Isaac-Cartpole-Albedo-Camera-Direct-v0", "cartpole"), + ("Isaac-Cartpole-Camera-Presets-Direct-v0", "cartpole"), + ("Isaac-Cartpole-Depth-Camera-Direct-v0", "cartpole"), + ("Isaac-Cartpole-RGB-Camera-Direct-v0", "cartpole"), + ("Isaac-Cartpole-SimpleShading-Constant-Camera-Direct-v0", "cartpole"), + ("Isaac-Cartpole-SimpleShading-Diffuse-Camera-Direct-v0", "cartpole"), + ("Isaac-Cartpole-SimpleShading-Full-Camera-Direct-v0", "cartpole"), + ("Isaac-Repose-Cube-Shadow-Vision-Direct-v0", "shadow_hand"), ] -@pytest.mark.parametrize("task_id", _RENDER_CORRECTNESS_TASK_IDS) -def test_registered_tasks(task_id): +@pytest.mark.parametrize("task_id, env_name", _RENDER_CORRECTNESS_TASK_IDS) +def test_registered_tasks(task_id, env_name): """Camera output must be non-empty for each registered task with camera-based observations.""" env = None try: env_cfg = parse_env_cfg(task_id, num_envs=4) - env_cfg.seed = _ENV_SEED env = gym.make(task_id, cfg=env_cfg) unwrapped: Any = env.unwrapped @@ -795,8 +795,6 @@ def test_registered_tasks(task_id): if sim is not None: sim._app_control_on_stop_handle = None - env.reset(seed=_ENV_SEED) - camera_outputs_nested_dict = _collect_camera_outputs(env) num_camera_outputs = len(camera_outputs_nested_dict) assert num_camera_outputs == 1, f"[{task_id}] Expected 1 camera output, got {num_camera_outputs}." @@ -808,7 +806,7 @@ def test_registered_tasks(task_id): "default_physics", "default_renderer", camera_outputs, - max_different_pixels_percentage=5.0, + max_different_pixels_percentage=_MAX_DIFFERENT_PIXELS_PERCENTAGE_BY_ENV_NAME[env_name], ) finally: if env is not None: diff --git a/source/isaaclab_tasks/test/test_sim_launcher_visualizer_intent.py b/source/isaaclab_tasks/test/test_sim_launcher_visualizer_intent.py index 853a9fb31a5d..c6bad5c19f1e 100644 --- a/source/isaaclab_tasks/test/test_sim_launcher_visualizer_intent.py +++ b/source/isaaclab_tasks/test/test_sim_launcher_visualizer_intent.py @@ -75,15 +75,18 @@ def set_bool(self, path: str, value: bool) -> None: monkeypatch.setattr( sim_launcher, "compute_kit_requirements", lambda env_cfg, launcher_args: (False, False, {"none"}) ) - monkeypatch.setitem( - sys.modules, - "isaaclab.app.settings_manager", - types.SimpleNamespace(get_settings_manager=lambda: _FakeSettings()), - ) + # `app_launcher` imports both names from settings_manager; provide a full stub module + # so `from isaaclab.app import AppLauncher` succeeds in kitless mode. + _sm = types.ModuleType("isaaclab.app.settings_manager") + _sm.get_settings_manager = lambda: _FakeSettings() + _sm.initialize_carb_settings = lambda: None + monkeypatch.setitem(sys.modules, "isaaclab.app.settings_manager", _sm) env_cfg = _DummyEnvCfg(_DummySimCfg(None)) launcher_args = argparse.Namespace(visualizer=["none"]) with sim_launcher.launch_simulation(env_cfg, launcher_args): pass - assert captured == {"types": "", "explicit": True, "disable_all": True} + # `sync_visualizer_cli_settings_to_carb` uses ``" ".join(visualizer)`` → ``"none"`` for ``["none"]``, + # not an empty string (empty only when *visualizer* is missing/empty). + assert captured == {"types": "none", "explicit": True, "disable_all": True} diff --git a/source/isaaclab_teleop/config/extension.toml b/source/isaaclab_teleop/config/extension.toml index 881c57a52727..13c63e04ab99 100644 --- a/source/isaaclab_teleop/config/extension.toml +++ b/source/isaaclab_teleop/config/extension.toml @@ -1,6 +1,6 @@ [package] # Semantic Versioning is used: https://semver.org/ -version = "0.3.5" +version = "0.3.6" # Description title = "Isaac Lab Teleop" diff --git a/source/isaaclab_teleop/docs/CHANGELOG.rst b/source/isaaclab_teleop/docs/CHANGELOG.rst index d526500022c1..9ad0bf77ecd2 100644 --- a/source/isaaclab_teleop/docs/CHANGELOG.rst +++ b/source/isaaclab_teleop/docs/CHANGELOG.rst @@ -1,6 +1,51 @@ Changelog --------- +0.3.6 (2026-04-21) +~~~~~~~~~~~~~~~~~~~ + +Added +^^^^^ + +* Added :attr:`~isaaclab_teleop.IsaacTeleopCfg.control_channel_uuid` for + receiving teleop control commands (start/stop/reset) from the headset via + an OpenXR message channel. The channel is managed by TeleopCore's native + ``teleop_control_pipeline`` mechanism. + +* Added :class:`~isaaclab_teleop.teleop_message_processor.TeleopMessageProcessor` + retargeter that converts raw message-channel payloads into boolean control + signals for :class:`~isaacteleop.teleop_session_manager.DefaultTeleopStateManager`. + +* Added :func:`~isaaclab_teleop.poll_control_events` helper, + :class:`~isaaclab_teleop.ControlEvents` dataclass, and + :class:`~isaaclab_teleop.SupportsControlEvents` protocol for polling + start/stop/reset signals from any teleop device in a single call. + +* Added :attr:`~isaaclab_teleop.IsaacTeleopDevice.last_control_events` + property exposing the most recent control events from the message channel. + Control events are automatically bridged to legacy + :meth:`~isaaclab_teleop.IsaacTeleopDevice.add_callback` callbacks. + +Changed +^^^^^^^ + +* :meth:`~isaaclab_teleop.IsaacTeleopDevice.reset` now injects a + ``reset`` :class:`ExecutionEvents` into TeleopCore's ``ComputeContext`` + on the next pipeline step, resetting retargeter cross-step state. + Previously only the XR anchor was reset. + +Fixed +^^^^^ + +* Fixed ``record_demos.py`` not resetting the teleop device when a + success condition triggers an environment reset. Retargeters now + reinitialize their state on success-triggered resets. + +* Fixed shutdown hang caused by Kit's pre-shutdown callback calling + ``stop()`` while the simulation loop was still running. The callback + now uses the same graceful teardown path as the XR-disabled handler. + + 0.3.5 (2026-04-06) ~~~~~~~~~~~~~~~~~~~ diff --git a/source/isaaclab_teleop/isaaclab_teleop/__init__.pyi b/source/isaaclab_teleop/isaaclab_teleop/__init__.pyi index 655c7025cb0f..045f16f0c690 100644 --- a/source/isaaclab_teleop/isaaclab_teleop/__init__.pyi +++ b/source/isaaclab_teleop/isaaclab_teleop/__init__.pyi @@ -6,15 +6,20 @@ __all__ = [ "CLOUDXR_AVP_ENV", "CLOUDXR_JS_ENV", + "ControlEvents", "IsaacTeleopCfg", "IsaacTeleopDevice", - "create_isaac_teleop_device", - "XrAnchorSynchronizer", + "SupportsControlEvents", + "TELEOP_CONTROL_CHANNEL_UUID", "XrAnchorRotationMode", + "XrAnchorSynchronizer", "XrCfg", + "create_isaac_teleop_device", + "poll_control_events", "remove_camera_configs", ] +from .control_events import TELEOP_CONTROL_CHANNEL_UUID, ControlEvents, SupportsControlEvents, poll_control_events from .isaac_teleop_cfg import CLOUDXR_AVP_ENV, CLOUDXR_JS_ENV, IsaacTeleopCfg from .isaac_teleop_device import IsaacTeleopDevice, create_isaac_teleop_device from .xr_anchor_utils import XrAnchorSynchronizer diff --git a/source/isaaclab_teleop/isaaclab_teleop/command_handler.py b/source/isaaclab_teleop/isaaclab_teleop/command_handler.py index eb5fb38aeb44..7e999e638f5e 100644 --- a/source/isaaclab_teleop/isaaclab_teleop/command_handler.py +++ b/source/isaaclab_teleop/isaaclab_teleop/command_handler.py @@ -3,57 +3,30 @@ # # SPDX-License-Identifier: BSD-3-Clause -"""Teleop command handling for IsaacTeleop-based teleoperation.""" +"""Teleop command callback registry for IsaacTeleop-based teleoperation.""" from __future__ import annotations -import logging from collections.abc import Callable -from typing import Any - -import carb - -logger = logging.getLogger(__name__) class CommandHandler: - """Handles teleop command callbacks and XR message bus events. - - This class is responsible for: - - 1. Registering callbacks for teleop commands (START, STOP, RESET) - 2. Subscribing to the XR message bus for command events - 3. Dispatching callbacks when commands are received - - Teleop commands can be triggered via XR controller buttons or the - message bus. The handler normalizes command names (e.g. mapping - ``"R"`` to ``"RESET"``) and dispatches to registered callbacks. + """Lightweight callback registry for teleop commands. + + Scripts can register callbacks for ``START``, ``STOP``, and ``RESET`` + commands via :meth:`add_callback`. The callbacks are dispatched by + :meth:`fire` when the corresponding command is received. + + Note: + In the current architecture control signals arrive through + TeleopCore's ``teleop_control_pipeline`` and are consumed via + :func:`~isaaclab_teleop.poll_control_events`. This registry is + retained for backward compatibility with scripts that register + callbacks before the pipeline-based path was introduced. """ - TELEOP_COMMAND_EVENT_TYPE = "teleop_command" - - def __init__(self, xr_core: Any | None = None, on_reset: Callable[[], None] | None = None): - """Initialize the command handler. - - Args: - xr_core: The XRCore singleton, or ``None`` if XR is not available. - When provided, the handler subscribes to the message bus for - teleop command events. - on_reset: Optional hook called whenever a ``"reset"`` message-bus - event is received, *in addition to* the user's RESET callback. - This allows the device to perform internal reset actions (e.g. - resetting the XR anchor) without coupling the handler to the - anchor manager. - """ + def __init__(self) -> None: self._callbacks: dict[str, Callable] = {} - self._on_reset = on_reset - self._xr_core = xr_core - self._vc_subscription = None - - if self._xr_core is not None: - self._vc_subscription = self._xr_core.get_message_bus().create_subscription_to_pop_by_type( - carb.events.type_from_string(self.TELEOP_COMMAND_EVENT_TYPE), self._on_teleop_command - ) @property def callbacks(self) -> dict[str, Callable]: @@ -70,7 +43,6 @@ def add_callback(self, key: str, func: Callable) -> None: func: The function to call when the command is received. Should take no arguments. """ - # Map "R" to "RESET" for compatibility with existing scripts if key == "R": key = "RESET" self._callbacks[key] = func @@ -84,19 +56,5 @@ def fire(self, command: str) -> None: if command in self._callbacks: self._callbacks[command]() - def _on_teleop_command(self, event: carb.events.IEvent) -> None: - """Handle teleop command events from the message bus.""" - msg = event.payload.get("message", "") - - if "start" in msg: - self.fire("START") - elif "stop" in msg: - self.fire("STOP") - elif "reset" in msg: - self.fire("RESET") - if self._on_reset is not None: - self._on_reset() - def cleanup(self) -> None: - """Release event subscriptions.""" - self._vc_subscription = None + """Release resources (no-op; retained for API compatibility).""" diff --git a/source/isaaclab_teleop/isaaclab_teleop/control_events.py b/source/isaaclab_teleop/isaaclab_teleop/control_events.py new file mode 100644 index 000000000000..69ddd7c1ed21 --- /dev/null +++ b/source/isaaclab_teleop/isaaclab_teleop/control_events.py @@ -0,0 +1,78 @@ +# 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 + +"""Teleop control events dataclass, polling helper, and well-known channel UUID.""" + +from __future__ import annotations + +import dataclasses +import uuid +from typing import Protocol, runtime_checkable + +TELEOP_CONTROL_CHANNEL_UUID: bytes = uuid.uuid5(uuid.NAMESPACE_DNS, "teleop_command").bytes +"""Well-known 16-byte UUID for the teleop control message channel. + +Derived deterministically as ``uuid5(NAMESPACE_DNS, "teleop_command")`` +so that both the Isaac Lab server and the Quest client can independently +compute the same channel identifier from the string ``"teleop_command"``. + +Pass this value as :attr:`~isaaclab_teleop.IsaacTeleopCfg.control_channel_uuid` +when configuring a teleop session with message-channel-based control. +""" + + +@dataclasses.dataclass(frozen=True) +class ControlEvents: + """Result of :func:`poll_control_events`. + + Attributes: + is_active: ``True`` when the teleop state machine is in RUNNING, + ``False`` when PAUSED or STOPPED, or ``None`` when no control + channel is configured (callers should leave their own active + flag unchanged). + should_reset: ``True`` when a reset was triggered this frame. + """ + + is_active: bool | None = None + should_reset: bool = False + + +_NO_OP_EVENTS = ControlEvents() +"""Shared immutable sentinel returned when no control channel is active.""" + + +@runtime_checkable +class SupportsControlEvents(Protocol): + """Duck type for teleop devices that expose control events.""" + + @property + def last_control_events(self) -> ControlEvents: ... + + +def poll_control_events(teleop_interface: SupportsControlEvents | object) -> ControlEvents: + """Poll control events from any teleop interface. + + Safe to call with any device type (keyboard, spacemouse, etc.). + Devices that do not expose the message-channel protocol return + a no-op :class:`ControlEvents`. + + Args: + teleop_interface: The teleop device to poll. Devices implementing + :class:`SupportsControlEvents` provide full type safety; other + devices are handled gracefully via duck typing. + + Returns: + A :class:`ControlEvents` with the latest start/stop and reset + signals. + """ + events = getattr(teleop_interface, "last_control_events", None) + if events is None: + return _NO_OP_EVENTS + if isinstance(events, ControlEvents): + return events + return ControlEvents( + is_active=getattr(events, "is_active", None), + should_reset=getattr(events, "should_reset", False), + ) diff --git a/source/isaaclab_teleop/isaaclab_teleop/isaac_teleop_cfg.py b/source/isaaclab_teleop/isaaclab_teleop/isaac_teleop_cfg.py index 6539fa67f346..f94a63d57589 100644 --- a/source/isaaclab_teleop/isaaclab_teleop/isaac_teleop_cfg.py +++ b/source/isaaclab_teleop/isaaclab_teleop/isaac_teleop_cfg.py @@ -14,6 +14,7 @@ from isaaclab.utils import configclass +from .control_events import TELEOP_CONTROL_CHANNEL_UUID from .xr_cfg import XrCfg _CLOUDXR_ENV_DIR = Path(__file__).resolve().parent @@ -117,6 +118,24 @@ def build_pipeline(): If ``None``, the tuning UI will not be opened. """ + control_channel_uuid: bytes | None = TELEOP_CONTROL_CHANNEL_UUID + """16-byte UUID for the teleop control message channel. + + Defaults to :data:`~isaaclab_teleop.TELEOP_CONTROL_CHANNEL_UUID` + (``uuid5(NAMESPACE_DNS, "teleop_command")``), which is the well-known + channel both the Isaac Lab server and CloudXR JS client use to + exchange start/stop/reset commands. + + When set, a ``teleop_control_pipeline`` is created automatically + using :class:`~isaaclab_teleop.teleop_message_processor.TeleopMessageProcessor` + and :class:`~isaacteleop.teleop_session_manager.DefaultTeleopStateManager`. + The remote client sends UTF-8 control commands over the OpenXR opaque + data channel identified by this UUID, and the results are exposed via + :func:`~isaaclab_teleop.poll_control_events`. + + Set to ``None`` to disable the control channel entirely. + """ + target_frame_prim_path: str | None = None """Optional USD prim path whose world frame becomes the target coordinate frame for all output poses. diff --git a/source/isaaclab_teleop/isaaclab_teleop/isaac_teleop_device.py b/source/isaaclab_teleop/isaaclab_teleop/isaac_teleop_device.py index 2e7c2c7a406b..3f8c565a7e21 100644 --- a/source/isaaclab_teleop/isaaclab_teleop/isaac_teleop_device.py +++ b/source/isaaclab_teleop/isaaclab_teleop/isaac_teleop_device.py @@ -15,6 +15,7 @@ import torch from .command_handler import CommandHandler +from .control_events import ControlEvents from .isaac_teleop_cfg import IsaacTeleopCfg from .session_lifecycle import TeleopSessionLifecycle from .xr_anchor_manager import XrAnchorManager @@ -35,8 +36,8 @@ class IsaacTeleopDevice: and coordinate-frame transform computation. * :class:`TeleopSessionLifecycle` -- pipeline building, OpenXR handle acquisition, session creation/destruction, and action-tensor extraction. - * :class:`CommandHandler` -- callback registration and XR message-bus - command dispatch. + * :class:`CommandHandler` -- callback registration for START / STOP / RESET + commands, bridged from the pipeline-based control events. Together they manage: @@ -67,7 +68,8 @@ class IsaacTeleopDevice: Teleop commands: The device supports callbacks for START, STOP, and RESET commands - that can be triggered via XR controller buttons or the message bus. + that can be triggered via the message-channel control pipeline or + registered directly via :meth:`add_callback`. Example: .. code-block:: python @@ -118,20 +120,16 @@ def __init__( """ self._cfg = cfg - # Compose the three collaborators self._anchor_manager = XrAnchorManager(cfg.xr_cfg) + self._command_handler = CommandHandler() self._session_lifecycle = TeleopSessionLifecycle( cfg, cloudxr_env_file=cloudxr_env_file, auto_launch_cloudxr=auto_launch_cloudxr, ) - self._command_handler = CommandHandler( - xr_core=self._anchor_manager.xr_core, - on_reset=self._anchor_manager.reset, - ) - # Controller button polling state (edge detection for right 'A') self._prev_right_a_pressed = False + self._prev_control_is_active: bool | None = None def __del__(self): """Clean up resources when the object is destroyed.""" @@ -188,9 +186,23 @@ def __exit__(self, exc_type, exc_val, exc_tb): def reset(self) -> None: """Reset the device state. - Resets the XR anchor synchronizer if present. + Resets the XR anchor synchronizer and schedules a + ``reset`` :class:`~isaacteleop.retargeting_engine.interface.execution_events.ExecutionEvents` + for the next pipeline step so that all retargeters reinitialize + their cross-step state. """ self._anchor_manager.reset() + self._session_lifecycle.request_reset() + + @property + def last_control_events(self) -> ControlEvents: + """Control events from the most recent :meth:`advance`. + + Returns a :class:`ControlEvents` derived from the teleop control + pipeline. When no control channel is configured, returns a + default (no-op) :class:`ControlEvents`. + """ + return self._session_lifecycle.last_control_events def add_callback(self, key: str, func: Callable) -> None: """Add a callback function for teleop commands. @@ -252,8 +264,39 @@ def advance(self, target_T_world: np.ndarray | torch.Tensor | SupportsDLPack | N # Poll controller buttons (e.g. toggle anchor rotation on right 'A' press) self._poll_buttons() + self._dispatch_control_callbacks() + return action + # ------------------------------------------------------------------ + # Control event -> callback bridge + # ------------------------------------------------------------------ + + def _dispatch_control_callbacks(self) -> None: + """Fire legacy callbacks when control events indicate a state change. + + This bridges the pipeline-based :class:`ControlEvents` with the + callback-based :class:`CommandHandler` so that scripts which registered + callbacks via :meth:`add_callback` still receive dispatches. + + Only fires START/STOP when ``is_active`` transitions between ``True`` + and ``False``; initial transitions from ``None`` are ignored to avoid + spurious callbacks during ``DefaultTeleopStateManager``'s + STOPPED -> PAUSED progression. + """ + from .control_events import _NO_OP_EVENTS + + events = self._session_lifecycle.last_control_events + if events is _NO_OP_EVENTS: + return + if events.should_reset: + self._command_handler.fire("RESET") + self._anchor_manager.reset() + if events.is_active is not None: + if self._prev_control_is_active is not None and events.is_active != self._prev_control_is_active: + self._command_handler.fire("START" if events.is_active else "STOP") + self._prev_control_is_active = events.is_active + # ------------------------------------------------------------------ # Target frame transform (config-driven rebase) # ------------------------------------------------------------------ diff --git a/source/isaaclab_teleop/isaaclab_teleop/session_lifecycle.py b/source/isaaclab_teleop/isaaclab_teleop/session_lifecycle.py index d395fab31a30..fa5f36f658e0 100644 --- a/source/isaaclab_teleop/isaaclab_teleop/session_lifecycle.py +++ b/source/isaaclab_teleop/isaaclab_teleop/session_lifecycle.py @@ -18,10 +18,13 @@ if TYPE_CHECKING: from isaacteleop.cloudxr import CloudXRLauncher from isaacteleop.oxr import OpenXRSessionHandles + from isaacteleop.retargeting_engine.interface.execution_events import ExecutionEvents from isaacteleop.retargeting_engine_ui import MultiRetargeterTuningUIImGui from isaacteleop.teleop_session_manager import TeleopSession +from .control_events import _NO_OP_EVENTS, ControlEvents from .isaac_teleop_cfg import IsaacTeleopCfg +from .teleop_message_processor import TeleopMessageProcessor class SupportsDLPack(Protocol): @@ -71,6 +74,19 @@ def _to_numpy_4x4(mat: np.ndarray | torch.Tensor | SupportsDLPack) -> np.ndarray return np.asarray(mat, dtype=np.float32) +def _execution_events_to_control(ee: ExecutionEvents) -> ControlEvents: + """Map TeleopCore :class:`ExecutionEvents` to the script-facing :class:`ControlEvents`.""" + from isaacteleop.retargeting_engine.interface.execution_events import ExecutionState + + if ee.execution_state == ExecutionState.RUNNING: + is_active: bool | None = True + elif ee.execution_state in (ExecutionState.PAUSED, ExecutionState.STOPPED): + is_active = False + else: + is_active = None + return ControlEvents(is_active=is_active, should_reset=ee.reset) + + class TeleopSessionLifecycle: """Manages the IsaacTeleop session lifecycle. @@ -78,11 +94,13 @@ class TeleopSessionLifecycle: 1. Building the retargeting pipeline from configuration 2. Adding a parallel ``ControllersSource`` for button-state access - 3. Acquiring OpenXR handles from Kit's XR bridge extension - 4. Creating, entering, and exiting the ``TeleopSession`` - 5. Building external inputs for pipeline leaf nodes (e.g. world-to-anchor transform) - 6. Stepping the session and extracting the flattened action tensor - 7. Managing the optional retargeting tuning UI + 3. Building the optional ``teleop_control_pipeline`` for headset-driven + start/stop/reset via a message channel + 4. Acquiring OpenXR handles from Kit's XR bridge extension + 5. Creating, entering, and exiting the ``TeleopSession`` + 6. Building external inputs for pipeline leaf nodes (e.g. world-to-anchor transform) + 7. Stepping the session and extracting the flattened action tensor + 8. Managing the optional retargeting tuning UI """ WORLD_T_ANCHOR_INPUT_NAME = "world_T_anchor" @@ -118,8 +136,12 @@ def __init__( # Session state (populated during start) self._session: TeleopSession | None = None self._pipeline = None + self._teleop_control_pipeline = None + self._message_processor: TeleopMessageProcessor | None = None self._last_right_controller = None self._session_start_deferred_logged = False + # Fallback for host-initiated resets when no control pipeline is configured + self._pending_reset = False # CloudXR runtime launcher (created in start if configured, stopped in stop) self._cloudxr_launcher: CloudXRLauncher | None = None @@ -192,6 +214,48 @@ def last_right_controller(self): """ return self._last_right_controller + @property + def has_control_channel(self) -> bool: + """Whether a message-channel-based control pipeline is configured.""" + return self._message_processor is not None + + @property + def last_control_events(self) -> ControlEvents: + """Control events from the most recent :meth:`step`. + + When a ``teleop_control_pipeline`` is configured, derives + :class:`ControlEvents` from + ``session.last_context.execution_events``. Otherwise returns a + default (no-op) :class:`ControlEvents`. + """ + if self._message_processor is None: + return _NO_OP_EVENTS + if self._session is None: + return _NO_OP_EVENTS + ctx = self._session.last_context + if ctx is None: + return _NO_OP_EVENTS + return _execution_events_to_control(ctx.execution_events) + + def request_reset(self) -> None: + """Schedule a reset for the next pipeline step. + + When a control pipeline is configured, the reset flows through + :meth:`TeleopMessageProcessor.inject_reset` so + :class:`~isaacteleop.teleop_session_manager.DefaultTeleopStateManager` + processes it normally. Otherwise falls back to an + ``execution_events`` override on the next :meth:`step` call. + + If the control channel already processed a reset this frame, + this method is a no-op to avoid a redundant second reset pulse. + """ + if self.last_control_events.should_reset: + return + if self._message_processor is not None: + self._message_processor.inject_reset() + else: + self._pending_reset = True + # ------------------------------------------------------------------ # Lifecycle: start / stop # ------------------------------------------------------------------ @@ -203,9 +267,10 @@ def start(self) -> None: the CloudXR runtime and WSS proxy are launched first. Builds the retargeting pipeline, wraps it with a parallel - ``ControllersSource`` for button-state access, attempts to acquire - OpenXR handles, and opens the retargeting tuning UI if retargeters - are configured. + ``ControllersSource`` for button-state access, builds the optional + ``teleop_control_pipeline`` for message-channel control, attempts + to acquire OpenXR handles, and opens the retargeting tuning UI if + retargeters are configured. If the OpenXR handles are not yet available (e.g. user hasn't clicked "Start AR"), session creation is deferred and will be retried on each @@ -222,12 +287,19 @@ def start(self) -> None: self._last_right_controller = None button_controllers = ControllersSource("_button_controllers") - self._pipeline = OutputCombiner( - { - "action": user_pipeline.output("action"), - self._CONTROLLER_RIGHT_KEY: button_controllers.output(ControllersSource.RIGHT), - } - ) + pipeline_outputs: dict[str, Any] = { + "action": user_pipeline.output("action"), + self._CONTROLLER_RIGHT_KEY: button_controllers.output(ControllersSource.RIGHT), + } + self._pipeline = OutputCombiner(pipeline_outputs) + + # Build optional teleop_control_pipeline for message-channel control + self._teleop_control_pipeline = None + self._message_processor = None + if self._cfg.control_channel_uuid is not None: + self._teleop_control_pipeline, self._message_processor = self._build_control_pipeline( + self._cfg.control_channel_uuid + ) # Try to start the session now; it may be deferred self._try_start_session() @@ -269,7 +341,12 @@ def stop(self, exc_type=None, exc_val=None, exc_tb=None) -> None: # expected and safe to suppress. logger.debug(f"Suppressed error during IsaacTeleop session cleanup: {e}") self._session = None - self._pipeline = None + + # Always clear pipeline state (session may never have been created if + # OpenXR handles were never available). + self._pipeline = None + self._teleop_control_pipeline = None + self._message_processor = None if self._cloudxr_launcher is not None: try: @@ -282,18 +359,68 @@ def stop(self, exc_type=None, exc_val=None, exc_tb=None) -> None: logger.info("IsaacTeleop session ended") + # ------------------------------------------------------------------ + # Control pipeline construction + # ------------------------------------------------------------------ + + @staticmethod + def _build_control_pipeline(channel_uuid: bytes) -> tuple[Any, TeleopMessageProcessor]: + """Build a ``teleop_control_pipeline`` from a message channel UUID. + + Wires ``MessageChannelSource`` -> :class:`TeleopMessageProcessor` + -> :class:`~isaacteleop.teleop_session_manager.DefaultTeleopStateManager`. + + Args: + channel_uuid: 16-byte UUID for the OpenXR opaque data channel. + + Returns: + A ``(teleop_control_pipeline, message_processor)`` tuple. + """ + from isaacteleop.retargeting_engine.deviceio_source_nodes import message_channel_config + from isaacteleop.teleop_session_manager import DefaultTeleopStateManager + + source, _sink = message_channel_config( + name="_teleop_control", + channel_uuid=channel_uuid, + ) + + processor = TeleopMessageProcessor(name="_teleop_msg_processor") + processor_graph = processor.connect({processor.INPUT_MESSAGES: source.output("messages_tracked")}) + + state_manager = DefaultTeleopStateManager(name="_teleop_state") + teleop_control_pipeline = state_manager.connect( + { + state_manager.INPUT_KILL: processor_graph.output("kill"), + state_manager.INPUT_RUN_TOGGLE: processor_graph.output("run_toggle"), + state_manager.INPUT_RESET: processor_graph.output("reset"), + } + ) + + return teleop_control_pipeline, processor + + # ------------------------------------------------------------------ + # Extension / XR lifecycle callbacks + # ------------------------------------------------------------------ + def _on_request_required_extensions(self) -> list[str]: """Callback for required extensions subscription. + Inspects both the main pipeline and the ``teleop_control_pipeline`` + (if configured) so that extensions required by the control channel + (e.g. ``XR_NV_opaque_data_channel``) are included. + Returns: A list of required extensions. """ from isaacteleop.teleop_session_manager.helpers import get_required_oxr_extensions_from_pipeline - required_extensions = ( - get_required_oxr_extensions_from_pipeline(self._pipeline) if self._pipeline is not None else [] - ) + required_extensions: list[str] = [] + if self._pipeline is not None: + required_extensions.extend(get_required_oxr_extensions_from_pipeline(self._pipeline)) + if self._teleop_control_pipeline is not None: + required_extensions.extend(get_required_oxr_extensions_from_pipeline(self._teleop_control_pipeline)) + required_extensions = sorted(set(required_extensions)) logger.info(f"Required extensions: {required_extensions}") return required_extensions @@ -307,10 +434,16 @@ def _on_xr_enabled_changed(self, item, event_type): self._teardown_dead_session() def _on_pre_shutdown(self, _event): - """Called when Kit is closing; run full cleanup since the app is exiting.""" + """Called when Kit is closing; tear down the session but leave the + pipeline intact so the main loop can exit via its own control flow + (``simulation_app.is_running()`` will go ``False``). + + Full resource cleanup happens later when the context manager's + ``__exit__`` calls :meth:`stop`. + """ logger.info("Shutting down IsaacTeleop session due to Kit close") self._pre_shutdown_subscription = None - self.stop() + self._teardown_dead_session() # ------------------------------------------------------------------ # Deferred session creation @@ -341,11 +474,6 @@ def _try_start_session(self) -> bool: if self._session is not None: return True - # In headless mode the AR profile setting is deliberately omitted - # from the .kit file so that all extensions (including the teleop - # bridge and its BridgeComponent) can load and register before Kit - # creates the OpenXR instance. We enable it here, after extensions - # are loaded; Kit will process the change on the next event-loop tick. self._ensure_xr_ar_profile_enabled() from isaacteleop.oxr import OpenXRSessionHandles @@ -371,6 +499,7 @@ def _try_start_session(self) -> bool: app_name=self._cfg.app_name, trackers=[], pipeline=self._pipeline, + teleop_control_pipeline=self._teleop_control_pipeline, plugins=self._cfg.plugins, oxr_handles=oxr_handles, ) @@ -436,6 +565,15 @@ def step( # pipeline contains ValueInput leaf nodes. external_inputs = self._build_external_inputs(anchor_world_matrix_fn, target_T_world) + # When no control pipeline is configured, host-initiated resets use + # the execution_events override as a fallback path. + execution_events = None + if self._pending_reset: + from isaacteleop.retargeting_engine.interface.execution_events import ExecutionEvents, ExecutionState + + execution_events = ExecutionEvents(reset=True, execution_state=ExecutionState.RUNNING) + self._pending_reset = False + # Execute one step of the teleop session. # If the underlying OpenXR session was destroyed externally (e.g. # user clicked "Stop AR"), the step call will fail. We catch the @@ -443,7 +581,10 @@ def step( # can continue rendering (or wait for the session to restart). assert self._session is not None # guaranteed by _try_start_session above try: - result = self._session.step(external_inputs=external_inputs) + result = self._session.step( + external_inputs=external_inputs, + execution_events=execution_events, + ) except Exception as e: logger.warning(f"IsaacTeleop session step failed (XR session likely torn down): {e}") self._teardown_dead_session() diff --git a/source/isaaclab_teleop/isaaclab_teleop/teleop_message_processor.py b/source/isaaclab_teleop/isaaclab_teleop/teleop_message_processor.py new file mode 100644 index 000000000000..1844925c4d0c --- /dev/null +++ b/source/isaaclab_teleop/isaaclab_teleop/teleop_message_processor.py @@ -0,0 +1,232 @@ +# 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 + +"""Message-channel payload parser for TeleopCore's teleop_control_pipeline. + +Provides :class:`TeleopMessageProcessor`, a lightweight +:class:`~isaacteleop.retargeting_engine.interface.BaseRetargeter` that +converts message-channel payloads into boolean pulse signals suitable for +:class:`~isaacteleop.teleop_session_manager.DefaultTeleopStateManager`. +""" + +from __future__ import annotations + +import json +import re +from typing import TYPE_CHECKING + +from isaacteleop.retargeting_engine.interface import BaseRetargeter, RetargeterIOType + +if TYPE_CHECKING: + from isaacteleop.retargeting_engine.interface.retargeter_core_types import ComputeContext, RetargeterIO + +_COMMAND_PATTERNS: list[tuple[re.Pattern[str], str]] = [ + (re.compile(r"\breset\b", re.IGNORECASE), "reset"), + (re.compile(r"\bstop\b", re.IGNORECASE), "stop"), + (re.compile(r"\bstart\b", re.IGNORECASE), "start"), +] +"""Ordered patterns for classifying a command string. + +``reset`` is checked first so that a hypothetical payload containing +both "reset" and "start" is treated as a reset (the more destructive +operation wins). ``stop`` precedes ``start`` for the same reason. +""" + +# Shadow states mirroring DefaultTeleopStateManager's ExecutionState. +_STOPPED = "stopped" +_PAUSED = "paused" +_RUNNING = "running" + +# DefaultTeleopStateManager cycles states on run_toggle rising edges: +# STOPPED -> PAUSED -> RUNNING -> PAUSED -> RUNNING -> ... +# To map imperative "start" (= go to RUNNING) and "stop" (= go to PAUSED) +# we emit the right number of toggle edges based on predicted state. +_START_TOGGLE_SEQUENCES: dict[str, list[bool]] = { + _STOPPED: [True, False, True], # 2 edges: STOPPED -> PAUSED -> RUNNING + _PAUSED: [True], # 1 edge: PAUSED -> RUNNING + _RUNNING: [], # already running +} +_STOP_TOGGLE_SEQUENCES: dict[str, list[bool]] = { + _RUNNING: [True], # 1 edge: RUNNING -> PAUSED + _PAUSED: [], # already paused + _STOPPED: [], # already stopped +} +# Shadow state advances on each rising edge (True after False). +_TOGGLE_TRANSITIONS: dict[str, str] = { + _STOPPED: _PAUSED, + _PAUSED: _RUNNING, + _RUNNING: _PAUSED, +} + + +class TeleopMessageProcessor(BaseRetargeter): + """Parse message-channel payloads into boolean control signals. + + Consumes the ``messages_tracked`` output of a + :class:`~isaacteleop.retargeting_engine.deviceio_source_nodes.MessageChannelSource` + and produces three boolean pulse outputs that drive + :class:`~isaacteleop.teleop_session_manager.DefaultTeleopStateManager`: + + * ``run_toggle`` -- pulsed ``True`` on rising edges; the number of + edges depends on the target state (e.g. ``"start"`` from STOPPED + emits two edges over three frames: STOPPED -> PAUSED -> RUNNING). + * ``kill`` -- always ``False`` (reserved for fail-safe; ``"stop"`` + uses ``run_toggle`` to reach PAUSED instead of STOPPED). + * ``reset`` -- pulsed ``True`` for one frame on ``"reset"``. + + The processor maintains a *shadow state* that mirrors + ``DefaultTeleopStateManager``'s internal state so it can emit the + correct toggle sequence for imperative commands. + + Payload formats supported: + + 1. **JSON (Quest client format)**:: + + {"type": "teleop_command", "message": {"command": "start teleop"}} + + 2. **Plain text (fallback)**: raw UTF-8 string matched by word boundary + (``"start"``, ``"stop"``, ``"reset"``). + + Host-initiated resets (e.g. environment success) are injected via + :meth:`inject_reset`, which sets the ``reset`` output ``True`` on the + next compute call without requiring a message-channel payload. + """ + + INPUT_MESSAGES = "messages_tracked" + + def __init__(self, name: str) -> None: + self._inject_reset_pending = False + self._shadow_state = _STOPPED + self._run_toggle_queue: list[bool] = [] + self._prev_toggle_output = False + super().__init__(name=name) + + def inject_reset(self) -> None: + """Schedule a reset pulse on the next pipeline step. + + The ``reset`` output will be ``True`` for exactly one frame, then + automatically cleared. + """ + self._inject_reset_pending = True + + def _make_toggle_sequence(self, base_sequence: list[bool]) -> list[bool]: + """Prepend a ``False`` frame if needed to guarantee a clean rising edge. + + ``DefaultTeleopStateManager`` uses edge detection + (``pressed and not prev_pressed``), so emitting ``True`` when the + previous output was already ``True`` would not trigger a state + transition. This method prepends ``False`` when necessary. + """ + if not base_sequence: + return [] + seq = list(base_sequence) + if self._prev_toggle_output: + seq.insert(0, False) + return seq + + def input_spec(self) -> RetargeterIOType: + from isaacteleop.retargeting_engine.deviceio_source_nodes.deviceio_tensor_types import ( + MessageChannelMessagesTrackedGroup, + ) + + return {self.INPUT_MESSAGES: MessageChannelMessagesTrackedGroup()} + + def output_spec(self) -> RetargeterIOType: + from isaacteleop.teleop_session_manager.teleop_state_manager_types import bool_signal + + return { + "run_toggle": bool_signal("run_toggle"), + "kill": bool_signal("kill"), + "reset": bool_signal("reset"), + } + + def _compute_fn( + self, + inputs: RetargeterIO, + outputs: RetargeterIO, + context: ComputeContext, + ) -> None: + del context + + reset = self._inject_reset_pending + self._inject_reset_pending = False + + # Parse incoming messages and enqueue toggle sequences. + messages_tracked = inputs[self.INPUT_MESSAGES][0] + data = getattr(messages_tracked, "data", None) + if data: + for message in data: + payload = getattr(message, "payload", None) + if payload is None: + continue + try: + text = bytes(payload).decode("utf-8") + except (UnicodeDecodeError, TypeError): + continue + + command = _extract_command(text) + if command is None: + continue + + kind = _classify_command(command) + if kind == "start" and not self._run_toggle_queue: + self._run_toggle_queue = self._make_toggle_sequence(_START_TOGGLE_SEQUENCES[self._shadow_state]) + elif kind == "stop" and not self._run_toggle_queue: + self._run_toggle_queue = self._make_toggle_sequence(_STOP_TOGGLE_SEQUENCES[self._shadow_state]) + elif kind == "reset": + reset = True + + # Drain the toggle queue (one value per frame). + if self._run_toggle_queue: + run_toggle = self._run_toggle_queue.pop(0) + else: + run_toggle = False + + # Advance shadow state on rising edges (matches DefaultTeleopStateManager's + # edge detection: ``pressed and not prev_pressed``). + if run_toggle and not self._prev_toggle_output: + self._shadow_state = _TOGGLE_TRANSITIONS[self._shadow_state] + self._prev_toggle_output = run_toggle + + outputs["run_toggle"][0] = run_toggle + outputs["kill"][0] = False + outputs["reset"][0] = reset + + +def _classify_command(text: str) -> str | None: + """Return ``"start"``, ``"stop"``, ``"reset"``, or ``None``. + + Uses word-boundary matching so that e.g. ``"stop_and_restart"`` + matches ``"stop"`` (not ``"start"``). + """ + for pattern, label in _COMMAND_PATTERNS: + if pattern.search(text): + return label + return None + + +def _extract_command(text: str) -> str | None: + """Extract the command string from a JSON or plain-text payload. + + Tries JSON parsing first (Quest client format) and falls back to the + raw text for plain-string payloads. Non-string JSON scalars (numbers, + arrays, booleans) are discarded. + """ + try: + obj = json.loads(text) + except (json.JSONDecodeError, TypeError): + return text + + if not isinstance(obj, dict): + return None + if obj.get("type") != "teleop_command": + return None + + msg = obj.get("message") + if isinstance(msg, dict): + return msg.get("command", "") + if isinstance(msg, str): + return msg + return None diff --git a/source/isaaclab_teleop/test/test_cloudxr_lifecycle.py b/source/isaaclab_teleop/test/test_cloudxr_lifecycle.py index e9565a7d3e41..43131f70cfc3 100644 --- a/source/isaaclab_teleop/test/test_cloudxr_lifecycle.py +++ b/source/isaaclab_teleop/test/test_cloudxr_lifecycle.py @@ -38,8 +38,15 @@ "isaacteleop.oxr", "isaacteleop.retargeting_engine", "isaacteleop.retargeting_engine.interface", + "isaacteleop.retargeting_engine.interface.execution_events", + "isaacteleop.retargeting_engine.interface.retargeter_core_types", + "isaacteleop.retargeting_engine.interface.tensor_group_type", + "isaacteleop.retargeting_engine.deviceio_source_nodes", + "isaacteleop.retargeting_engine.deviceio_source_nodes.deviceio_tensor_types", "isaacteleop.retargeting_engine_ui", "isaacteleop.teleop_session_manager", + "isaacteleop.teleop_session_manager.teleop_state_manager_retargeter", + "isaacteleop.teleop_session_manager.teleop_state_manager_types", "isaacsim", "isaacsim.kit", "isaacsim.kit.xr", @@ -85,6 +92,7 @@ def _make_cfg() -> IsaacTeleopCfg: """Build a minimal IsaacTeleopCfg with a dummy pipeline_builder.""" return IsaacTeleopCfg( pipeline_builder=lambda: MagicMock(), + control_channel_uuid=None, ) diff --git a/source/isaaclab_teleop/test/test_control_events.py b/source/isaaclab_teleop/test/test_control_events.py new file mode 100644 index 000000000000..8bc05d3f957c --- /dev/null +++ b/source/isaaclab_teleop/test/test_control_events.py @@ -0,0 +1,586 @@ +# 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 + +# pyright: reportPrivateUsage=none + +"""Tests for TeleopMessageProcessor, _classify_command, _extract_command, +and poll_control_events. + +These tests exercise pure logic (no Omniverse/Isaac Sim stack required). +The message processor is tested by calling its ``_compute_fn`` method +directly with fake pipeline I/O, mirroring how TeleopCore's +``teleop_control_pipeline`` mechanism invokes it. +""" + +from __future__ import annotations + +import dataclasses +import json +import sys +from types import ModuleType +from unittest.mock import MagicMock + +import pytest + +# --------------------------------------------------------------------------- +# Stub out isaacteleop modules before any isaaclab_teleop imports so the +# tests can run in a plain Python environment without Omniverse. +# --------------------------------------------------------------------------- + +_MODULES_TO_STUB = [ + "isaacteleop", + "isaacteleop.deviceio", + "isaacteleop.deviceio_trackers", + "isaacteleop.retargeting_engine", + "isaacteleop.retargeting_engine.deviceio_source_nodes", + "isaacteleop.retargeting_engine.deviceio_source_nodes.deviceio_tensor_types", + "isaacteleop.retargeting_engine.interface", + "isaacteleop.retargeting_engine.interface.retargeter_core_types", + "isaacteleop.retargeting_engine.interface.tensor_group_type", + "isaacteleop.retargeting_engine_ui", + "isaacteleop.schema", + "isaacteleop.teleop_session_manager", + "isaacteleop.teleop_session_manager.teleop_state_manager_retargeter", + "isaacteleop.teleop_session_manager.teleop_state_manager_types", +] + +_stubs: dict[str, ModuleType | MagicMock] = {} + + +def _install_stubs(): + for name in _MODULES_TO_STUB: + if name not in sys.modules: + _stubs[name] = MagicMock() + sys.modules[name] = _stubs[name] + + from enum import Enum + + class ExecutionState(str, Enum): + UNKNOWN = "unknown" + STOPPED = "stopped" + PAUSED = "paused" + RUNNING = "running" + + @dataclasses.dataclass + class ExecutionEvents: + reset: bool = False + execution_state: ExecutionState = ExecutionState.UNKNOWN + + ee_mod = sys.modules["isaacteleop.retargeting_engine.interface.execution_events"] = ModuleType( + "isaacteleop.retargeting_engine.interface.execution_events" + ) + ee_mod.ExecutionState = ExecutionState # type: ignore[attr-defined] + ee_mod.ExecutionEvents = ExecutionEvents # type: ignore[attr-defined] + + iface = sys.modules["isaacteleop.retargeting_engine.interface"] + iface.ExecutionState = ExecutionState # type: ignore[attr-defined] + iface.ExecutionEvents = ExecutionEvents # type: ignore[attr-defined] + iface.RetargeterIOType = dict # type: ignore[attr-defined] + + class FakeBaseRetargeter: + def __init__(self, name: str) -> None: + self.name = name + + iface.BaseRetargeter = FakeBaseRetargeter # type: ignore[attr-defined] + + tsm_types = sys.modules["isaacteleop.teleop_session_manager.teleop_state_manager_types"] + tsm_types.bool_signal = MagicMock # type: ignore[attr-defined] + + dt_mod = sys.modules["isaacteleop.retargeting_engine.deviceio_source_nodes.deviceio_tensor_types"] + dt_mod.MessageChannelMessagesTrackedGroup = MagicMock # type: ignore[attr-defined] + + +_install_stubs() + +from isaaclab_teleop.control_events import ControlEvents, poll_control_events # noqa: E402 +from isaaclab_teleop.teleop_message_processor import ( # noqa: E402 + TeleopMessageProcessor, + _classify_command, + _extract_command, +) + +# --------------------------------------------------------------------------- +# Test doubles for MessageChannelMessagesTrackedT +# --------------------------------------------------------------------------- + + +@dataclasses.dataclass +class _FakePayload: + payload: bytes + + +@dataclasses.dataclass +class _FakeTracked: + data: list[_FakePayload] | None = None + + +def _tracked(*payloads: bytes) -> _FakeTracked: + """Build a lightweight stand-in for ``MessageChannelMessagesTrackedT``.""" + return _FakeTracked(data=[_FakePayload(p) for p in payloads]) + + +def _empty_tracked() -> _FakeTracked: + return _FakeTracked(data=[]) + + +def _null_tracked() -> _FakeTracked: + return _FakeTracked(data=None) + + +def _make_inputs(messages_tracked): + """Build a fake RetargeterIO dict for the processor.""" + tg = MagicMock() + tg.__getitem__ = MagicMock(return_value=messages_tracked) + return {TeleopMessageProcessor.INPUT_MESSAGES: tg} + + +class _FakeOutputSlot: + """Captures ``outputs["key"][0] = value`` assignments.""" + + def __init__(self): + self.value = None + + def __setitem__(self, idx, val): + self.value = val + + def __getitem__(self, idx): + return self.value + + +def _make_outputs(): + """Build a fake outputs dict with capturable slots.""" + return {"run_toggle": _FakeOutputSlot(), "kill": _FakeOutputSlot(), "reset": _FakeOutputSlot()} + + +def _step(proc, messages_tracked) -> dict: + """Run the processor's _compute_fn and return captured outputs.""" + inputs = _make_inputs(messages_tracked) + outputs = _make_outputs() + proc._compute_fn(inputs, outputs, context=None) + return {k: v.value for k, v in outputs.items()} + + +# =========================================================================== +# TeleopMessageProcessor: basic command parsing +# =========================================================================== + + +class TestStartCommand: + def test_start_sets_run_toggle(self): + proc = TeleopMessageProcessor(name="test") + result = _step(proc, _tracked(b"start")) + assert result["run_toggle"] is True + assert result["kill"] is False + assert result["reset"] is False + + def test_start_does_not_set_reset(self): + proc = TeleopMessageProcessor(name="test") + result = _step(proc, _tracked(b"start")) + assert result["reset"] is False + + +class TestStopCommand: + def test_stop_from_stopped_is_noop(self): + proc = TeleopMessageProcessor(name="test") + result = _step(proc, _tracked(b"stop")) + assert result["run_toggle"] is False + assert result["kill"] is False + + +class TestResetCommand: + def test_reset_sets_reset_flag(self): + proc = TeleopMessageProcessor(name="test") + result = _step(proc, _tracked(b"reset")) + assert result["reset"] is True + assert result["run_toggle"] is False + assert result["kill"] is False + + +class TestResetPulseBehaviour: + def test_reset_clears_on_next_step(self): + proc = TeleopMessageProcessor(name="test") + result = _step(proc, _tracked(b"reset")) + assert result["reset"] is True + + result = _step(proc, _empty_tracked()) + assert result["reset"] is False + + +class TestKillAlwaysFalse: + def test_kill_is_always_false(self): + proc = TeleopMessageProcessor(name="test") + for payload in [b"start", b"stop", b"reset", b"hello"]: + result = _step(proc, _tracked(payload)) + assert result["kill"] is False + + +# =========================================================================== +# Shadow state and toggle sequences +# =========================================================================== + + +class TestStartFromStopped: + """``start`` from STOPPED needs 2 toggle edges over 3 frames.""" + + def test_full_sequence_reaches_running(self): + proc = TeleopMessageProcessor(name="test") + # Frame 0: "start" received, first toggle edge queued + r0 = _step(proc, _tracked(b"start")) + assert r0["run_toggle"] is True # edge 1: STOPPED -> PAUSED + + # Frame 1: queue drains False (prev resets) + r1 = _step(proc, _empty_tracked()) + assert r1["run_toggle"] is False + + # Frame 2: queue drains True (second edge) + r2 = _step(proc, _empty_tracked()) + assert r2["run_toggle"] is True # edge 2: PAUSED -> RUNNING + + # Frame 3: queue empty, back to idle + r3 = _step(proc, _empty_tracked()) + assert r3["run_toggle"] is False + + def test_shadow_state_is_running_after_sequence(self): + proc = TeleopMessageProcessor(name="test") + _step(proc, _tracked(b"start")) + _step(proc, _empty_tracked()) + _step(proc, _empty_tracked()) + assert proc._shadow_state == "running" + + +class TestStartFromPaused: + """``start`` from PAUSED needs 1 toggle edge.""" + + def test_single_edge_reaches_running(self): + proc = TeleopMessageProcessor(name="test") + # Drive to RUNNING: start sequence plays 3 frames + _step(proc, _tracked(b"start")) + _step(proc, _empty_tracked()) + _step(proc, _empty_tracked()) + assert proc._shadow_state == "running" + + # Stop to reach PAUSED (prev_toggle is True from start sequence, + # so a False is prepended before the toggle edge) + _step(proc, _tracked(b"stop")) # drains False (prepended) + r_stop_edge = _step(proc, _empty_tracked()) # drains True (edge) + assert r_stop_edge["run_toggle"] is True + assert proc._shadow_state == "paused" + + # Start from PAUSED: prev_toggle is True, so False prepended + _step(proc, _tracked(b"start")) # drains False (prepended) + r_start_edge = _step(proc, _empty_tracked()) # drains True (edge) + assert r_start_edge["run_toggle"] is True + assert proc._shadow_state == "running" + + +class TestStartFromRunning: + """``start`` when already RUNNING is a no-op.""" + + def test_start_from_running_noop(self): + proc = TeleopMessageProcessor(name="test") + _step(proc, _tracked(b"start")) + _step(proc, _empty_tracked()) + _step(proc, _empty_tracked()) + assert proc._shadow_state == "running" + + result = _step(proc, _tracked(b"start")) + assert result["run_toggle"] is False + + +class TestStopFromRunning: + """``stop`` from RUNNING uses one toggle edge to reach PAUSED.""" + + def test_stop_pauses(self): + proc = TeleopMessageProcessor(name="test") + _step(proc, _tracked(b"start")) + _step(proc, _empty_tracked()) + _step(proc, _empty_tracked()) + assert proc._shadow_state == "running" + + # prev_toggle is True, so stop prepends False before the edge + r0 = _step(proc, _tracked(b"stop")) + assert r0["run_toggle"] is False # prepended False + r1 = _step(proc, _empty_tracked()) + assert r1["run_toggle"] is True # edge: RUNNING -> PAUSED + assert proc._shadow_state == "paused" + + +class TestStopFromPaused: + """``stop`` when already PAUSED is a no-op.""" + + def test_stop_from_paused_noop(self): + proc = TeleopMessageProcessor(name="test") + _step(proc, _tracked(b"start")) + _step(proc, _empty_tracked()) + _step(proc, _empty_tracked()) + # Stop to PAUSED + _step(proc, _tracked(b"stop")) + _step(proc, _empty_tracked()) + assert proc._shadow_state == "paused" + + result = _step(proc, _tracked(b"stop")) + assert result["run_toggle"] is False + + +class TestCommandDuringToggleSequence: + """Commands received while a toggle sequence is in progress are ignored.""" + + def test_second_start_during_sequence_ignored(self): + proc = TeleopMessageProcessor(name="test") + _step(proc, _tracked(b"start")) # starts the 3-frame sequence + # Second start during the sequence should not restart it + r1 = _step(proc, _tracked(b"start")) + assert r1["run_toggle"] is False # draining the False from queue + + r2 = _step(proc, _empty_tracked()) + assert r2["run_toggle"] is True # second edge fires normally + + +# =========================================================================== +# inject_reset +# =========================================================================== + + +class TestInjectReset: + def test_inject_reset_produces_pulse(self): + proc = TeleopMessageProcessor(name="test") + proc.inject_reset() + result = _step(proc, _empty_tracked()) + assert result["reset"] is True + + def test_inject_reset_clears_after_one_step(self): + proc = TeleopMessageProcessor(name="test") + proc.inject_reset() + _step(proc, _empty_tracked()) + result = _step(proc, _empty_tracked()) + assert result["reset"] is False + + def test_inject_reset_combines_with_message_reset(self): + proc = TeleopMessageProcessor(name="test") + proc.inject_reset() + result = _step(proc, _tracked(b"reset")) + assert result["reset"] is True + + def test_inject_reset_independent_of_toggle(self): + proc = TeleopMessageProcessor(name="test") + proc.inject_reset() + result = _step(proc, _tracked(b"start")) + assert result["run_toggle"] is True + assert result["reset"] is True + + +# =========================================================================== +# Word boundary matching +# =========================================================================== + + +class TestWordBoundaryMatching: + @pytest.mark.parametrize("payload", [b"teleop start", b"xr start session", b"start now"]) + def test_start_word(self, payload: bytes): + proc = TeleopMessageProcessor(name="test") + result = _step(proc, _tracked(payload)) + assert result["run_toggle"] is True + + @pytest.mark.parametrize("payload", [b"teleop reset", b"env reset"]) + def test_reset_word(self, payload: bytes): + proc = TeleopMessageProcessor(name="test") + result = _step(proc, _tracked(payload)) + assert result["reset"] is True + + +class TestAmbiguousPayloads: + def test_reset_wins_over_start(self): + proc = TeleopMessageProcessor(name="test") + result = _step(proc, _tracked(b"reset and start")) + assert result["reset"] is True + assert result["run_toggle"] is False + + +# =========================================================================== +# Empty, null, and malformed batches +# =========================================================================== + + +class TestEmptyAndNullBatches: + def test_empty_data_list(self): + proc = TeleopMessageProcessor(name="test") + result = _step(proc, _empty_tracked()) + assert result["run_toggle"] is False + assert result["kill"] is False + assert result["reset"] is False + + def test_null_data(self): + proc = TeleopMessageProcessor(name="test") + result = _step(proc, _null_tracked()) + assert result["run_toggle"] is False + + def test_none_input(self): + proc = TeleopMessageProcessor(name="test") + result = _step(proc, None) + assert result["run_toggle"] is False + + +class TestMultipleMessagesInBatch: + def test_start_then_reset_in_one_batch(self): + proc = TeleopMessageProcessor(name="test") + result = _step(proc, _tracked(b"start", b"reset")) + assert result["run_toggle"] is True + assert result["reset"] is True + + +class TestMalformedPayloads: + def test_invalid_utf8(self): + proc = TeleopMessageProcessor(name="test") + result = _step(proc, _tracked(b"\xff\xfe")) + assert result["run_toggle"] is False + assert result["kill"] is False + assert result["reset"] is False + + def test_none_payload(self): + proc = TeleopMessageProcessor(name="test") + tracked = _FakeTracked(data=[_FakePayload(payload=None)]) # type: ignore[arg-type] + result = _step(proc, tracked) + assert result["run_toggle"] is False + + +# =========================================================================== +# JSON format tests (Quest client sends JSON teleop_command messages) +# =========================================================================== + + +def _json_command(command: str) -> bytes: + """Build a Quest-style JSON teleop_command payload.""" + return json.dumps({"type": "teleop_command", "message": {"command": command}}).encode("utf-8") + + +class TestJsonFormat: + def test_json_start_teleop(self): + proc = TeleopMessageProcessor(name="test") + result = _step(proc, _tracked(_json_command("start teleop"))) + assert result["run_toggle"] is True + + def test_json_stop_teleop_from_stopped_noop(self): + proc = TeleopMessageProcessor(name="test") + result = _step(proc, _tracked(_json_command("stop teleop"))) + assert result["run_toggle"] is False + + def test_json_reset_teleop(self): + proc = TeleopMessageProcessor(name="test") + result = _step(proc, _tracked(_json_command("reset teleop"))) + assert result["reset"] is True + + def test_json_wrong_type_ignored(self): + payload = json.dumps({"type": "other_event", "message": {"command": "start"}}).encode("utf-8") + proc = TeleopMessageProcessor(name="test") + result = _step(proc, _tracked(payload)) + assert result["run_toggle"] is False + + def test_json_message_as_string(self): + payload = json.dumps({"type": "teleop_command", "message": "start teleop"}).encode("utf-8") + proc = TeleopMessageProcessor(name="test") + result = _step(proc, _tracked(payload)) + assert result["run_toggle"] is True + + +# =========================================================================== +# _extract_command unit tests +# =========================================================================== + + +class TestExtractCommand: + def test_plain_text(self): + assert _extract_command("start teleop") == "start teleop" + + def test_json_teleop_command(self): + text = json.dumps({"type": "teleop_command", "message": {"command": "stop"}}) + assert _extract_command(text) == "stop" + + def test_json_wrong_type(self): + text = json.dumps({"type": "other", "message": {"command": "start"}}) + assert _extract_command(text) is None + + def test_json_no_message_key(self): + text = json.dumps({"type": "teleop_command"}) + assert _extract_command(text) is None + + def test_json_non_dict_value_returns_none(self): + assert _extract_command("42") is None + assert _extract_command("[1, 2, 3]") is None + assert _extract_command("true") is None + + +# =========================================================================== +# _classify_command unit tests +# =========================================================================== + + +class TestClassifyCommand: + def test_exact_words(self): + assert _classify_command("start") == "start" + assert _classify_command("stop") == "stop" + assert _classify_command("reset") == "reset" + + def test_word_boundary_prevents_false_match(self): + assert _classify_command("upstart") is None + assert _classify_command("nonstop") is None + assert _classify_command("unreset") is None + + def test_reset_beats_start(self): + assert _classify_command("reset and start") == "reset" + + def test_stop_beats_start(self): + assert _classify_command("stop and start") == "stop" + + def test_unrecognized_text(self): + assert _classify_command("hello world") is None + + def test_case_insensitive(self): + assert _classify_command("START") == "start" + assert _classify_command("Stop Teleop") == "stop" + assert _classify_command("RESET NOW") == "reset" + + +# =========================================================================== +# poll_control_events tests +# =========================================================================== + + +class TestPollControlEvents: + def test_plain_object_returns_default(self): + result = poll_control_events(object()) + assert result.is_active is None + assert result.should_reset is False + + def test_device_with_control_events(self): + class FakeDevice: + @property + def last_control_events(self): + return ControlEvents(is_active=True, should_reset=True) + + result = poll_control_events(FakeDevice()) + assert result.is_active is True + assert result.should_reset is True + + def test_device_with_none_events(self): + class FakeDevice: + last_control_events = None + + result = poll_control_events(FakeDevice()) + assert result.is_active is None + assert result.should_reset is False + + def test_duck_typed_snapshot(self): + class FakeSnapshot: + is_active = False + should_reset = True + + class FakeDevice: + @property + def last_control_events(self): + return FakeSnapshot() + + result = poll_control_events(FakeDevice()) + assert result.is_active is False + assert result.should_reset is True 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() diff --git a/source/isaaclab_visualizers/isaaclab_visualizers/kit/kit_visualizer.py b/source/isaaclab_visualizers/isaaclab_visualizers/kit/kit_visualizer.py index 6b1b5c2077dc..ca5117f62d09 100644 --- a/source/isaaclab_visualizers/isaaclab_visualizers/kit/kit_visualizer.py +++ b/source/isaaclab_visualizers/isaaclab_visualizers/kit/kit_visualizer.py @@ -11,10 +11,13 @@ import logging from typing import TYPE_CHECKING -from pxr import UsdGeom +from pxr import Usd, UsdGeom, Vt +from isaaclab.app.settings_manager import get_settings_manager from isaaclab.visualizers.base_visualizer import BaseVisualizer +from isaaclab_visualizers.newton_adapter import resolve_visible_env_indices + from .kit_visualizer_cfg import KitVisualizerCfg logger = logging.getLogger(__name__) @@ -22,6 +25,8 @@ if TYPE_CHECKING: from isaaclab.physics import BaseSceneDataProvider +_DEFAULT_VIEWPORT_NAME = "Visualizer Viewport" + class KitVisualizer(BaseVisualizer): """Kit visualizer using Isaac Sim viewport.""" @@ -42,10 +47,11 @@ def __init__(self, cfg: KitVisualizerCfg): self._sim_time = 0.0 self._step_counter = 0 self._hidden_env_visibilities: dict[str, str] = {} - # Camera prim path that set_camera_view() writes to. Pinned at initialization so that - # user-switching the GUI viewport to a sensor camera does not corrupt the sensor's prim. - self._controlled_camera_path: str | None = None + # PointInstancer prim path -> (had authored invisibleIds, previous value) for partial viz restore. + self._point_instancer_invisible_ids_backup: dict[str, tuple[bool, object]] = {} self._runtime_headless = bool(cfg.headless) + # USD path for the viewport's active camera, refreshed after setup (used by CI/tests). + self._controlled_camera_path: str | None = None # ---- Lifecycle ------------------------------------------------------------------------ @@ -68,22 +74,29 @@ def initialize(self, scene_data_provider: BaseSceneDataProvider) -> None: metadata = scene_data_provider.get_metadata() self._ensure_simulation_app() - self._setup_viewport(usd_stage) + self._setup_viewport() self._env_ids = self._compute_visualized_env_ids() - if self._env_ids: + num_envs_meta = int(metadata.get("num_envs", 0)) + self._resolved_visible_env_ids = resolve_visible_env_indices( + self._env_ids, self.cfg.max_visible_envs, num_envs_meta + ) + if self._resolved_visible_env_ids is not None: logger.warning( - "[KitVisualizer] env_filter_ids filtering is cosmetic only (no perf gain) in OV; hiding other envs." + "[KitVisualizer] Partial visualization in Kit uses visibility only; unselected env prims are hidden." ) - self._apply_env_visibility(usd_stage, metadata) - num_visualized_envs = len(self._env_ids) if self._env_ids is not None else int(metadata.get("num_envs", 0)) + self._apply_env_visibility(usd_stage, metadata, self._resolved_visible_env_ids) + num_visualized_envs = ( + len(self._resolved_visible_env_ids) if self._resolved_visible_env_ids is not None else num_envs_meta + ) self._log_initialization_table( logger=logger, title="KitVisualizer Configuration", rows=[ - ("camera_position", self.cfg.camera_position), - ("camera_target", self.cfg.camera_target), - ("camera_source", self.cfg.camera_source), + ("eye", self.cfg.eye), + ("lookat", self.cfg.lookat), + ("cam_source", self.cfg.cam_source), + ("max_visible_envs", self.cfg.max_visible_envs), ("num_visualized_envs", num_visualized_envs), ("create_viewport", self.cfg.create_viewport), ("headless", self._runtime_headless), @@ -105,18 +118,19 @@ def step(self, dt: float) -> None: try: import omni.kit.app - from isaaclab.app.settings_manager import get_settings_manager - app = omni.kit.app.get_app() if app is not None and app.is_running(): - # Keep app pumping for viewport/UI updates only. - # Simulation stepping is owned by SimulationContext. + # Keep app pumping for viewport/UI updates only; physics is owned by SimulationContext. + # Disable playSimulations around app.update() so Kit does not advance its own physics here. settings = get_settings_manager() settings.set_bool("/app/player/playSimulations", False) app.update() settings.set_bool("/app/player/playSimulations", True) except (ImportError, AttributeError) as exc: logger.debug("[KitVisualizer] App update skipped: %s", exc) + # Markers (VisualizationMarkers) are often created or resized to num_envs only after the first + # simulation / debug-vis step; re-apply PointInstancer invisibleIds each step when partial viz is on. + self._refresh_partial_viz_point_instancers_if_needed() def close(self) -> None: """Close viewport resources and restore temporary state.""" @@ -150,8 +164,6 @@ def is_running(self) -> bool: def is_training_paused(self) -> bool: """Return whether simulation play flag is paused in Kit settings.""" try: - from isaaclab.app.settings_manager import get_settings_manager - settings = get_settings_manager() play_flag = settings.get("/app/player/playSimulations") return play_flag is False @@ -179,10 +191,17 @@ def set_camera_view( ) -> None: """Set active viewport camera eye/target. + When :attr:`self.cfg.cam_source` is ``"cfg"``, this is a no-op: the pose comes only from + :attr:`self.cfg.eye` / :attr:`self.cfg.lookat` (applied in :meth:`_setup_viewport`). Otherwise + :class:`~isaaclab.sim.simulation_context.SimulationContext` and :class:`ViewportCameraController` + would overwrite that pose with :class:`~isaaclab.envs.common.ViewerCfg`-driven views. + Args: eye: Camera eye position. target: Camera look-at target. """ + if self.cfg.cam_source == "cfg": + return if not self._is_initialized: logger.debug("[KitVisualizer] set_camera_view() ignored because visualizer is not initialized.") return @@ -215,22 +234,33 @@ def _ensure_simulation_app(self) -> None: except ImportError: pass - def _setup_viewport(self, usd_stage) -> None: - """Create/resolve viewport and configure initial camera. - - Args: - usd_stage: USD stage used for camera prim setup. - """ + def _setup_viewport(self) -> None: + """Create/resolve viewport and configure initial camera.""" import omni.kit.viewport.utility as vp_utils from omni.ui import DockPosition if self._runtime_headless: - # In headless mode we keep the visualizer active but skip viewport/window setup. + # Headless: no viewport window; apply cfg pose to the default perspective camera path. self._viewport_window = None self._viewport_api = None + if self.cfg.cam_source == "prim_path": + logger.warning( + "[KitVisualizer] cam_source='prim_path' has limited support in headless mode; " + "using eye/lookat from cfg instead." + ) + self._apply_cfg_camera_pose_if_configured() + self._refresh_controlled_camera_path() return - if self.cfg.create_viewport and self.cfg.viewport_name: + effective_viewport_name = ( + self.cfg.viewport_name if self.cfg.viewport_name is not None else _DEFAULT_VIEWPORT_NAME + ) + + if self.cfg.create_viewport: + if not str(effective_viewport_name).strip(): + raise RuntimeError( + "[KitVisualizer] viewport_name must be a non-empty string when create_viewport=True." + ) dock_position_name = self.cfg.dock_position.upper() dock_position_map = { "LEFT": DockPosition.LEFT, @@ -241,7 +271,7 @@ def _setup_viewport(self, usd_stage) -> None: dock_pos = dock_position_map.get(dock_position_name, DockPosition.SAME) self._viewport_window = vp_utils.create_viewport_window( - name=self.cfg.viewport_name, + name=effective_viewport_name, width=self.cfg.window_width, height=self.cfg.window_height, position_x=50, @@ -249,28 +279,33 @@ def _setup_viewport(self, usd_stage) -> None: docked=True, ) - asyncio.ensure_future(self._dock_viewport_async(self.cfg.viewport_name, dock_pos)) - self._create_and_assign_camera(usd_stage) + asyncio.ensure_future(self._dock_viewport_async(effective_viewport_name, dock_pos)) else: self._viewport_window = vp_utils.get_active_viewport_window() if self._viewport_window is None: logger.warning("[KitVisualizer] No active viewport window found.") self._viewport_api = None + self._refresh_controlled_camera_path() return self._viewport_api = self._viewport_window.viewport_api - # Pin the camera path we will write to, using the active camera at init time. - # This must happen before any _set_viewport_camera() call so the path is known. - self._controlled_camera_path = self._viewport_api.get_active_camera() or "/OmniverseKit_Persp" - if self.cfg.camera_source == "usd_path": - if not self._set_active_camera_path(self.cfg.camera_usd_path): - logger.warning( - "[KitVisualizer] camera_usd_path '%s' not found; using configured camera.", - self.cfg.camera_usd_path, + if self.cfg.cam_source == "prim_path": + if not self._set_active_camera_path(self.cfg.cam_prim_path): + raise RuntimeError( + "[KitVisualizer] cam_source='prim_path' requires a valid cam_prim_path. " + f"Camera prim not found: '{self.cfg.cam_prim_path}'." ) - self._set_viewport_camera(self.cfg.camera_position, self.cfg.camera_target) else: - self._set_viewport_camera(self.cfg.camera_position, self.cfg.camera_target) + self._apply_cfg_camera_pose_if_configured() + self._refresh_controlled_camera_path() + + def _refresh_controlled_camera_path(self) -> None: + """Cache :attr:`_controlled_camera_path` from the active viewport (or default persp).""" + if self._viewport_api is not None: + path = self._viewport_api.get_active_camera() + self._controlled_camera_path = path if path else "/OmniverseKit_Persp" + else: + self._controlled_camera_path = "/OmniverseKit_Persp" async def _dock_viewport_async(self, viewport_name: str, dock_position) -> None: """Dock a created viewport window relative to main viewport.""" @@ -303,35 +338,23 @@ async def _dock_viewport_async(self, viewport_name: str, dock_position) -> None: await omni.kit.app.get_app().next_update_async() viewport_window.focus() - def _create_and_assign_camera(self, usd_stage) -> None: - """Create viewport camera prim (if needed) and set it active.""" - camera_path = f"/World/Cameras/{self.cfg.viewport_name}_Camera".replace(" ", "_") - - camera_prim = usd_stage.GetPrimAtPath(camera_path) - if not camera_prim.IsValid(): - UsdGeom.Camera.Define(usd_stage, camera_path) - - if self._viewport_api: - self._viewport_api.set_active_camera(camera_path) - self._controlled_camera_path = camera_path - def _set_viewport_camera(self, position: tuple[float, float, float], target: tuple[float, float, float]) -> None: """Apply eye/target camera view to the active viewport.""" import isaacsim.core.utils.viewports as isaacsim_viewports - if self._viewport_api is None: - return - # Use the camera path pinned at initialization. This prevents user-switching the GUI - # viewport to a sensor camera from corrupting the sensor's USD prim transform. - camera_path = self._controlled_camera_path - if not camera_path: - camera_path = self._viewport_api.get_active_camera() if self._viewport_api else None + camera_path = None + if self._viewport_api is not None: + camera_path = self._viewport_api.get_active_camera() if not camera_path: camera_path = "/OmniverseKit_Persp" + kwargs = {"eye": list(position), "target": list(target), "camera_prim_path": camera_path} + if self._viewport_api is not None: + kwargs["viewport_api"] = self._viewport_api + isaacsim_viewports.set_camera_view(**kwargs) - isaacsim_viewports.set_camera_view( - eye=list(position), target=list(target), camera_prim_path=camera_path, viewport_api=self._viewport_api - ) + def _apply_cfg_camera_pose_if_configured(self) -> None: + """Apply configured camera pose from eye/lookat.""" + self._set_viewport_camera(self.cfg.eye, self.cfg.lookat) def _set_active_camera_path(self, camera_path: str) -> bool: """Set active camera path for viewport if the prim exists. @@ -348,17 +371,14 @@ def _set_active_camera_path(self, camera_path: str) -> bool: if not camera_prim.IsValid(): return False self._viewport_api.set_active_camera(camera_path) - self._controlled_camera_path = camera_path return True - def _apply_env_visibility(self, usd_stage, metadata: dict) -> None: - """Hide non-selected environments for cosmetic env filtering.""" - if not self._env_ids: - return + def _apply_env_visibility(self, usd_stage, metadata: dict, visible_env_ids: list[int]) -> None: + """Hide environments not listed in ``visible_env_ids`` (cosmetic partial visualization).""" num_envs = int(metadata.get("num_envs", 0)) if num_envs <= 0: return - visible = set(self._env_ids) + visible = set(visible_env_ids) for env_id in range(num_envs): if env_id in visible: continue @@ -375,10 +395,63 @@ def _apply_env_visibility(self, usd_stage, metadata: dict) -> None: self._hidden_env_visibilities[env_path] = prev attr.Set(UsdGeom.Tokens.invisible) - def _restore_env_visibility(self) -> None: - """Restore environment visibilities modified by env filtering.""" - if not self._hidden_env_visibilities: + self._apply_visual_point_instancer_visibility(usd_stage, num_envs, visible) + + def _refresh_partial_viz_point_instancers_if_needed(self) -> None: + """Re-apply ``invisibleIds`` for env-scaled `/Visuals` instancers (handles lazy marker creation).""" + if self._resolved_visible_env_ids is None or self._scene_data_provider is None: return + usd_stage = self._scene_data_provider.get_usd_stage() + if usd_stage is None: + return + num_envs = int(self._scene_data_provider.get_metadata().get("num_envs", 0)) + if num_envs <= 0: + return + self._apply_visual_point_instancer_visibility(usd_stage, num_envs, set(self._resolved_visible_env_ids)) + + def _apply_visual_point_instancer_visibility(self, usd_stage, num_envs: int, visible_env_ids: set[int]) -> None: + """Set ``PointInstancer.invisibleIds`` for per-env `/Visuals` markers (e.g. velocity arrows).""" + hidden = [i for i in range(num_envs) if i not in visible_env_ids] + vt_hidden = Vt.Int64Array([int(i) for i in hidden]) + for root_path in ("/Visuals", "/World/Visuals"): + root_prim = usd_stage.GetPrimAtPath(root_path) + if not root_prim.IsValid(): + continue + for prim in Usd.PrimRange(root_prim): + if not prim.IsA(UsdGeom.PointInstancer): + continue + pi = UsdGeom.PointInstancer(prim) + n = self._point_instancer_instance_count(pi) + if n is None or n != num_envs: + continue + path_str = prim.GetPath().pathString + inv_attr = pi.GetInvisibleIdsAttr() + # Record original authorship/value once per instancer for :meth:`_restore_env_visibility`. + if path_str not in self._point_instancer_invisible_ids_backup: + was_authored = inv_attr.HasAuthoredValue() + prev = inv_attr.Get() if was_authored else None + self._point_instancer_invisible_ids_backup[path_str] = (was_authored, prev) + inv_attr.Set(vt_hidden) + + @staticmethod + def _point_instancer_instance_count(pi: UsdGeom.PointInstancer) -> int | None: + """Return instance count from the first authored per-instance array, if any.""" + for attr in ( + pi.GetPositionsAttr(), + pi.GetScalesAttr(), + pi.GetOrientationsAttr(), + pi.GetProtoIndicesAttr(), + ): + if not attr.HasAuthoredValue(): + continue + val = attr.Get() + if val is None: + continue + return len(val) + return None + + def _restore_env_visibility(self) -> None: + """Restore environment visibilities and PointInstancer ``invisibleIds`` from partial viz.""" usd_stage = self._scene_data_provider.get_usd_stage() if self._scene_data_provider else None if usd_stage is None: return @@ -391,3 +464,14 @@ def _restore_env_visibility(self) -> None: continue imageable.GetVisibilityAttr().Set(prev) self._hidden_env_visibilities.clear() + + for path_str, (was_authored, prev) in self._point_instancer_invisible_ids_backup.items(): + prim = usd_stage.GetPrimAtPath(path_str) + if not prim.IsValid() or not prim.IsA(UsdGeom.PointInstancer): + continue + inv_attr = UsdGeom.PointInstancer(prim).GetInvisibleIdsAttr() + if not was_authored: + inv_attr.Clear() + else: + inv_attr.Set(prev) + self._point_instancer_invisible_ids_backup.clear() diff --git a/source/isaaclab_visualizers/isaaclab_visualizers/kit/kit_visualizer_cfg.py b/source/isaaclab_visualizers/isaaclab_visualizers/kit/kit_visualizer_cfg.py index 88112a6f20b4..342be3fc2c6f 100644 --- a/source/isaaclab_visualizers/isaaclab_visualizers/kit/kit_visualizer_cfg.py +++ b/source/isaaclab_visualizers/isaaclab_visualizers/kit/kit_visualizer_cfg.py @@ -5,6 +5,8 @@ """Configuration for Kit-based visualizer.""" +from __future__ import annotations + from isaaclab.utils import configclass from isaaclab.visualizers.visualizer_cfg import VisualizerCfg @@ -16,20 +18,23 @@ class KitVisualizerCfg(VisualizerCfg): visualizer_type: str = "kit" """Type identifier for Kit visualizer.""" - viewport_name: str | None = "Visualizer Viewport" - """Viewport name to use. If None, uses active viewport.""" + viewport_name: str | None = None + """Name for a new viewport window when :attr:`create_viewport` is ``True``. + + If ``None``, a default name (``"Visualizer Viewport"``) is used. + """ create_viewport: bool = False - """Create new viewport with specified name and camera pose.""" + """If ``True``, create a new viewport window; if ``False``, use the active viewport window.""" headless: bool = False """Run without creating viewport windows when supported by the app.""" dock_position: str = "SAME" - """Dock position for new viewport. Options: 'LEFT', 'RIGHT', 'BOTTOM', 'SAME'.""" + """Dock position for a new viewport. Options: 'LEFT', 'RIGHT', 'BOTTOM', 'SAME'.""" window_width: int = 1280 - """Viewport width in pixels.""" + """Viewport width in pixels (when :attr:`create_viewport` is ``True``).""" window_height: int = 720 - """Viewport height in pixels.""" + """Viewport height in pixels (when :attr:`create_viewport` is ``True``).""" diff --git a/source/isaaclab_visualizers/isaaclab_visualizers/newton/newton_visualizer.py b/source/isaaclab_visualizers/isaaclab_visualizers/newton/newton_visualizer.py index 9ffeb062065d..8c8bb0bed9d8 100644 --- a/source/isaaclab_visualizers/isaaclab_visualizers/newton/newton_visualizer.py +++ b/source/isaaclab_visualizers/isaaclab_visualizers/newton/newton_visualizer.py @@ -16,6 +16,8 @@ from isaaclab.visualizers.base_visualizer import BaseVisualizer +from isaaclab_visualizers.newton_adapter import apply_viewer_visible_worlds, resolve_visible_env_indices + from .newton_visualizer_cfg import NewtonVisualizerCfg logger = logging.getLogger(__name__) @@ -56,7 +58,7 @@ def __init__( self._fallback_draw_controls = True def is_training_paused(self) -> bool: - """Return whether training is paused by viewer controls.""" + """Return whether simulation is paused by viewer controls.""" return self._paused_training def is_rendering_paused(self) -> bool: @@ -68,7 +70,7 @@ def _render_training_controls(self, imgui): imgui.separator() imgui.text("IsaacLab Controls") - pause_label = "Resume Training" if self._paused_training else "Pause Training" + pause_label = "Resume Simulation" if self._paused_training else "Pause Simulation" if imgui.button(pause_label): self._paused_training = not self._paused_training @@ -110,7 +112,7 @@ def _render_ui(self): imgui.set_next_window_pos(imgui.ImVec2(320, 10)) flags = 0 - if imgui.begin("Training Controls", flags=flags): + if imgui.begin("Simulation Controls", flags=flags): self._render_training_controls(imgui) imgui.end() return None @@ -281,40 +283,37 @@ def initialize(self, scene_data_provider: BaseSceneDataProvider) -> None: self._scene_data_provider = scene_data_provider metadata = scene_data_provider.get_metadata() + num_envs = int(metadata.get("num_envs", 0)) self._env_ids = self._compute_visualized_env_ids() - if self._env_ids: - get_filtered_model = getattr(scene_data_provider, "get_newton_model_for_env_ids", None) - if callable(get_filtered_model): - self._model = get_filtered_model(self._env_ids) - else: - self._model = scene_data_provider.get_newton_model() - else: - self._model = scene_data_provider.get_newton_model() - self._state = scene_data_provider.get_newton_state(self._env_ids) - - try: - self._viewer = NewtonViewerGL( - width=self.cfg.window_width, - height=self.cfg.window_height, - headless=self.cfg.headless, - metadata=metadata, - update_frequency=self.cfg.update_frequency, - ) - except Exception as exc: - if not self.cfg.headless: - raise - self._viewer = None - self._headless_no_viewer = True - logger.info( - "[NewtonVisualizer] Headless fallback enabled (ViewerGL unavailable in this environment): %s", - exc, - ) + self._model = scene_data_provider.get_newton_model() + self._state = scene_data_provider.get_newton_state() + + # Use pyglet's EGL headless backend when requested. Must run before the first + # ``pyglet.window`` import so ``Window`` resolves to :class:`~pyglet.window.headless.HeadlessWindow`. + if self.cfg.headless: + import pyglet + + pyglet.options["headless"] = True + + self._viewer = NewtonViewerGL( + width=self.cfg.window_width, + height=self.cfg.window_height, + headless=self.cfg.headless, + metadata=metadata, + update_frequency=self.cfg.update_frequency, + ) if self._viewer is not None: - max_worlds = self.cfg.max_worlds - self._viewer.set_model(self._model, max_worlds=max_worlds) + self._viewer.set_model(self._model) + apply_viewer_visible_worlds( + self._viewer, + env_ids=self._env_ids, + max_visible_envs=self.cfg.max_visible_envs, + num_envs=num_envs, + ) self._viewer.set_world_offsets((0.0, 0.0, 0.0)) - self._apply_camera_pose(self._resolve_initial_camera_pose()) + initial_pose = self._resolve_initial_camera_pose() + self._apply_camera_pose(initial_pose) self._viewer.up_axis = 2 # Z-up self._viewer.scaling = 1.0 @@ -336,19 +335,18 @@ def initialize(self, scene_data_provider: BaseSceneDataProvider) -> None: self._viewer.renderer.sky_lower = self._viewer._coerce_color3(self.cfg.sky_lower_color) self._viewer.renderer._light_color = self._viewer._coerce_color3(self.cfg.light_color) - num_visualized_envs = len(self._env_ids) if self._env_ids is not None else int(metadata.get("num_envs", 0)) + _resolved = resolve_visible_env_indices(self._env_ids, self.cfg.max_visible_envs, num_envs) + num_visualized_envs = len(_resolved) if _resolved is not None else num_envs self._log_initialization_table( logger=logger, title="NewtonVisualizer Configuration", rows=[ ( - "camera_position", - tuple(float(x) for x in self._viewer.camera.pos) - if self._viewer is not None - else self.cfg.camera_position, + "eye", + tuple(float(x) for x in self._viewer.camera.pos) if self._viewer is not None else self.cfg.eye, ), - ("camera_target", self._last_camera_pose[1] if self._last_camera_pose else self.cfg.camera_target), - ("camera_source", self.cfg.camera_source), + ("lookat", self._last_camera_pose[1] if self._last_camera_pose else self.cfg.lookat), + ("cam_source", self.cfg.cam_source), ("num_visualized_envs", num_visualized_envs), ("headless", self.cfg.headless), ], @@ -369,13 +367,13 @@ def step(self, dt: float) -> None: if self._viewer is None: if self._scene_data_provider is not None: - self._state = self._scene_data_provider.get_newton_state(self._env_ids) + self._state = self._scene_data_provider.get_newton_state() return - if self.cfg.camera_source == "usd_path": + if self.cfg.cam_source == "prim_path": self._update_camera_from_usd_path() - self._state = self._scene_data_provider.get_newton_state(self._env_ids) + self._state = self._scene_data_provider.get_newton_state() contacts = None if self._viewer.show_contacts: @@ -437,15 +435,15 @@ def _resolve_initial_camera_pose(self) -> tuple[tuple[float, float, float], tupl Returns: Camera eye and target tuples. """ - if self.cfg.camera_source == "usd_path": - pose = self._resolve_camera_pose_from_usd_path(self.cfg.camera_usd_path) + if self.cfg.cam_source == "prim_path": + pose = self._resolve_camera_pose_from_usd_path(self.cfg.cam_prim_path) if pose is not None: return pose - logger.warning( - "[NewtonVisualizer] camera_usd_path '%s' not found; using configured camera.", - self.cfg.camera_usd_path, + raise RuntimeError( + "[NewtonVisualizer] cam_source='prim_path' requires a resolvable camera prim path, " + f"but no camera pose was found for '{self.cfg.cam_prim_path}'." ) - return self.cfg.camera_position, self.cfg.camera_target + return self._resolve_cfg_camera_pose("NewtonVisualizer") def _apply_camera_pose(self, pose: tuple[tuple[float, float, float], tuple[float, float, float]]) -> None: """Apply camera eye/target pose to the Newton viewer. @@ -469,7 +467,7 @@ def _apply_camera_pose(self, pose: tuple[tuple[float, float, float], tuple[float def _update_camera_from_usd_path(self) -> None: """Refresh camera pose from configured USD camera path when it changes.""" - pose = self._resolve_camera_pose_from_usd_path(self.cfg.camera_usd_path) + pose = self._resolve_camera_pose_from_usd_path(self.cfg.cam_prim_path) if pose is None: return if self._last_camera_pose == pose: diff --git a/source/isaaclab_visualizers/isaaclab_visualizers/newton/newton_visualizer_cfg.py b/source/isaaclab_visualizers/isaaclab_visualizers/newton/newton_visualizer_cfg.py index b89e0a2d547c..711e86e03b31 100644 --- a/source/isaaclab_visualizers/isaaclab_visualizers/newton/newton_visualizer_cfg.py +++ b/source/isaaclab_visualizers/isaaclab_visualizers/newton/newton_visualizer_cfg.py @@ -25,12 +25,6 @@ class NewtonVisualizerCfg(VisualizerCfg): headless: bool = False """Run the Newton viewer without requiring a display server.""" - max_worlds: int | None = None - """Maximum number of worlds/environments rendered by the viewer. - - Set to ``None`` to leave this option disabled. - """ - update_frequency: int = 1 """Visualizer update frequency (updates every N frames).""" diff --git a/source/isaaclab_visualizers/isaaclab_visualizers/newton_adapter.py b/source/isaaclab_visualizers/isaaclab_visualizers/newton_adapter.py new file mode 100644 index 000000000000..6bc3d5a2b4f1 --- /dev/null +++ b/source/isaaclab_visualizers/isaaclab_visualizers/newton_adapter.py @@ -0,0 +1,63 @@ +# 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 helpers for viewer env selection (Newton viewers and Kit partial USD visibility).""" + +from __future__ import annotations + + +def resolve_visible_env_indices( + env_ids: list[int] | None, + max_visible_envs: int | None, + num_envs: int, +) -> list[int] | None: + """Resolve which env indices stay visible (same rules as :func:`apply_viewer_visible_worlds`). + + * Cap-only path (``env_ids`` is ``None``): contiguous ``0 .. min(cap, num_envs) - 1`` when ``max_visible_envs`` + is set; otherwise ``None`` (viewer shows all worlds). (Random cap-only selection is applied earlier by + turning it into explicit ``env_ids``.) + * Explicit path (``env_ids`` is a list): if ``max_visible_envs`` is set, keep only the first *cap* indices + (truncate from the end); if ``None``, use the full list. + + Returns: + Selected indices, or ``None`` when all environments should be visible (cap-only, no limit). + """ + if env_ids is not None: + out = list(env_ids) + if max_visible_envs is not None: + out = out[: max(0, int(max_visible_envs))] + return out + if max_visible_envs is not None and num_envs > 0: + n = min(int(max_visible_envs), num_envs) + return list(range(n)) + return None + + +def apply_viewer_visible_worlds( + viewer, + *, + env_ids: list[int] | None, + max_visible_envs: int | None, + num_envs: int, +) -> None: + """Select which simulation worlds are visualized; no-op if the viewer does not support it. + + Prefer this over ``set_model(..., max_worlds=...)`` (deprecated in Newton). + + Args: + viewer: Newton viewer (ViewerGL, ViewerRerun, ViewerViser, etc.). + env_ids: Env indices from ``visible_env_indices`` (after validation), or ``None`` for the cap-only + contiguous path (see ``VisualizerCfg``). + max_visible_envs: When ``env_ids`` is ``None``, caps the contiguous count; otherwise truncates the list to + the first *N* indices. + num_envs: Total environment count from scene metadata. + """ + if not hasattr(viewer, "set_visible_worlds"): + return + resolved = resolve_visible_env_indices(env_ids, max_visible_envs, num_envs) + if resolved is None: + viewer.set_visible_worlds(None) + else: + viewer.set_visible_worlds(resolved) diff --git a/source/isaaclab_visualizers/isaaclab_visualizers/rerun/rerun_visualizer.py b/source/isaaclab_visualizers/isaaclab_visualizers/rerun/rerun_visualizer.py index ab2ed723223f..5390802df69d 100644 --- a/source/isaaclab_visualizers/isaaclab_visualizers/rerun/rerun_visualizer.py +++ b/source/isaaclab_visualizers/isaaclab_visualizers/rerun/rerun_visualizer.py @@ -20,6 +20,8 @@ from isaaclab.visualizers.base_visualizer import BaseVisualizer +from isaaclab_visualizers.newton_adapter import apply_viewer_visible_worlds, resolve_visible_env_indices + from .rerun_visualizer_cfg import RerunVisualizerCfg if TYPE_CHECKING: @@ -145,16 +147,10 @@ def initialize(self, scene_data_provider: BaseSceneDataProvider) -> None: self._scene_data_provider = scene_data_provider metadata = scene_data_provider.get_metadata() + num_envs = int(metadata.get("num_envs", 0)) self._env_ids = self._compute_visualized_env_ids() - if self._env_ids: - get_filtered_model = getattr(scene_data_provider, "get_newton_model_for_env_ids", None) - if callable(get_filtered_model): - self._model = get_filtered_model(self._env_ids) - else: - self._model = scene_data_provider.get_newton_model() - else: - self._model = scene_data_provider.get_newton_model() - self._state = scene_data_provider.get_newton_state(self._env_ids) + self._model = scene_data_provider.get_newton_model() + self._state = scene_data_provider.get_newton_state() grpc_port = int(self.cfg.grpc_port) web_port = int(self.cfg.web_port) @@ -185,22 +181,30 @@ def initialize(self, scene_data_provider: BaseSceneDataProvider) -> None: viewer_url = _rerun_web_viewer_url(viewer_host, web_port, rerun_address) if self.cfg.open_browser and not start_server_in_viewer: _open_rerun_web_viewer(viewer_host, web_port, rerun_address) - self._viewer.set_model(self._model, max_worlds=self.cfg.max_worlds) + self._viewer.set_model(self._model) + apply_viewer_visible_worlds( + self._viewer, + env_ids=self._env_ids, + max_visible_envs=self.cfg.max_visible_envs, + num_envs=num_envs, + ) # Preserve simulation world positions (env_spacing) rather than adding viewer-side offsets. self._viewer.set_world_offsets((0.0, 0.0, 0.0)) - self._apply_camera_pose(self._resolve_initial_camera_pose()) + initial_pose = self._resolve_initial_camera_pose() + self._apply_camera_pose(initial_pose) self._viewer.up_axis = 2 self._viewer.scaling = 1.0 self._viewer._paused = False - num_visualized_envs = len(self._env_ids) if self._env_ids is not None else int(metadata.get("num_envs", 0)) + _resolved = resolve_visible_env_indices(self._env_ids, self.cfg.max_visible_envs, num_envs) + num_visualized_envs = len(_resolved) if _resolved is not None else num_envs self._log_initialization_table( logger=logger, title="RerunVisualizer Configuration", rows=[ - ("camera_position", self.cfg.camera_position), - ("camera_target", self.cfg.camera_target), - ("camera_source", self.cfg.camera_source), + ("eye", self.cfg.eye), + ("lookat", self.cfg.lookat), + ("cam_source", self.cfg.cam_source), ("num_visualized_envs", num_visualized_envs), ("endpoint", f"http://{viewer_host}:{web_port}"), ("viewer_url", viewer_url), @@ -227,10 +231,10 @@ def step(self, dt: float) -> None: self._sim_time += dt self._step_counter += 1 - if self.cfg.camera_source == "usd_path": + if self.cfg.cam_source == "prim_path": self._update_camera_from_usd_path() - self._state = self._scene_data_provider.get_newton_state(self._env_ids) + self._state = self._scene_data_provider.get_newton_state() if not self._viewer.is_paused(): self._viewer.begin_frame(self._sim_time) @@ -275,11 +279,15 @@ def is_running(self) -> bool: def _resolve_initial_camera_pose(self) -> tuple[tuple[float, float, float], tuple[float, float, float]]: """Resolve initial camera pose from config or USD camera path.""" - if self.cfg.camera_source == "usd_path": - pose = self._resolve_camera_pose_from_usd_path(self.cfg.camera_usd_path) + if self.cfg.cam_source == "prim_path": + pose = self._resolve_camera_pose_from_usd_path(self.cfg.cam_prim_path) if pose is not None: return pose - return self.cfg.camera_position, self.cfg.camera_target + raise RuntimeError( + "[RerunVisualizer] cam_source='prim_path' requires a resolvable camera prim path, " + f"but no camera pose was found for '{self.cfg.cam_prim_path}'." + ) + return self._resolve_cfg_camera_pose("RerunVisualizer") def _apply_camera_pose(self, pose: tuple[tuple[float, float, float], tuple[float, float, float]]) -> None: """Apply camera pose to rerun's 3D view controls. @@ -307,7 +315,7 @@ def _apply_camera_pose(self, pose: tuple[tuple[float, float, float], tuple[float def _update_camera_from_usd_path(self) -> None: """Refresh camera pose from configured USD camera path when it changes.""" - pose = self._resolve_camera_pose_from_usd_path(self.cfg.camera_usd_path) + pose = self._resolve_camera_pose_from_usd_path(self.cfg.cam_prim_path) if pose is None: return if self._last_camera_pose == pose: diff --git a/source/isaaclab_visualizers/isaaclab_visualizers/rerun/rerun_visualizer_cfg.py b/source/isaaclab_visualizers/isaaclab_visualizers/rerun/rerun_visualizer_cfg.py index 5edd918929de..780b346f802b 100644 --- a/source/isaaclab_visualizers/isaaclab_visualizers/rerun/rerun_visualizer_cfg.py +++ b/source/isaaclab_visualizers/isaaclab_visualizers/rerun/rerun_visualizer_cfg.py @@ -47,9 +47,3 @@ class RerunVisualizerCfg(VisualizerCfg): record_to_rrd: str | None = None """Path to save .rrd recording file. None = no recording.""" - - max_worlds: int | None = None - """Maximum number of worlds/environments rendered by the viewer. - - Set to ``None`` to leave this option disabled. - """ diff --git a/source/isaaclab_visualizers/isaaclab_visualizers/viser/viser_visualizer.py b/source/isaaclab_visualizers/isaaclab_visualizers/viser/viser_visualizer.py index c20bfcd85a9f..a629ab8b2fed 100644 --- a/source/isaaclab_visualizers/isaaclab_visualizers/viser/viser_visualizer.py +++ b/source/isaaclab_visualizers/isaaclab_visualizers/viser/viser_visualizer.py @@ -19,6 +19,8 @@ from isaaclab.visualizers.base_visualizer import BaseVisualizer +from isaaclab_visualizers.newton_adapter import apply_viewer_visible_worlds, resolve_visible_env_indices + from .viser_visualizer_cfg import ViserVisualizerCfg logger = logging.getLogger(__name__) @@ -143,28 +145,22 @@ def initialize(self, scene_data_provider: BaseSceneDataProvider) -> None: self._scene_data_provider = scene_data_provider metadata = scene_data_provider.get_metadata() self._env_ids = self._compute_visualized_env_ids() - if self._env_ids: - get_filtered_model = getattr(scene_data_provider, "get_newton_model_for_env_ids", None) - self._model = ( - get_filtered_model(self._env_ids) - if callable(get_filtered_model) - else scene_data_provider.get_newton_model() - ) - else: - self._model = scene_data_provider.get_newton_model() - self._state = scene_data_provider.get_newton_state(self._env_ids) + self._model = scene_data_provider.get_newton_model() + self._state = scene_data_provider.get_newton_state() self._active_record_path = self.cfg.record_to_viser self._create_viewer(record_to_viser=self.cfg.record_to_viser, metadata=metadata) - num_visualized_envs = len(self._env_ids) if self._env_ids is not None else int(metadata.get("num_envs", 0)) + num_envs_meta = int(metadata.get("num_envs", 0)) + _resolved = resolve_visible_env_indices(self._env_ids, self.cfg.max_visible_envs, num_envs_meta) + num_visualized_envs = len(_resolved) if _resolved is not None else num_envs_meta viewer_url = _viser_web_viewer_url(self.cfg.port) self._log_initialization_table( logger=logger, title="ViserVisualizer Configuration", rows=[ - ("camera_position", self.cfg.camera_position), - ("camera_target", self.cfg.camera_target), - ("camera_source", self.cfg.camera_source), + ("eye", self.cfg.eye), + ("lookat", self.cfg.lookat), + ("cam_source", self.cfg.cam_source), ("num_visualized_envs", num_visualized_envs), ("port", self.cfg.port), ("viewer_url", viewer_url), @@ -182,11 +178,11 @@ def step(self, dt: float) -> None: if not self._is_initialized or self._viewer is None or self._scene_data_provider is None: return - if self.cfg.camera_source == "usd_path": + if self.cfg.cam_source == "prim_path": self._update_camera_from_usd_path() self._apply_pending_camera_pose() - self._state = self._scene_data_provider.get_newton_state(self._env_ids) + self._state = self._scene_data_provider.get_newton_state() self._sim_time += dt self._viewer.begin_frame(self._sim_time) self._viewer.log_state(self._state) @@ -253,13 +249,20 @@ def _create_viewer(self, record_to_viser: str | None, metadata: dict | None = No record_to_viser=record_to_viser, metadata=metadata or {}, ) - max_worlds = self.cfg.max_worlds - self._viewer.set_model(self._model, max_worlds=max_worlds) + num_envs = int((metadata or {}).get("num_envs", 0)) + self._viewer.set_model(self._model) + apply_viewer_visible_worlds( + self._viewer, + env_ids=self._env_ids, + max_visible_envs=self.cfg.max_visible_envs, + num_envs=num_envs, + ) # Preserve simulation world positions (env_spacing) rather than adding viewer-side offsets. self._viewer.set_world_offsets((0.0, 0.0, 0.0)) if self.cfg.open_browser: _open_viser_web_viewer(self.cfg.port) - self._set_viser_camera_view(self._resolve_initial_camera_pose()) + initial_pose = self._resolve_initial_camera_pose() + self._set_viser_camera_view(initial_pose) self._sim_time = 0.0 def _close_viewer(self, finalize_viser: bool = False) -> None: @@ -277,15 +280,15 @@ def _close_viewer(self, finalize_viser: bool = False) -> None: def _resolve_initial_camera_pose(self) -> tuple[tuple[float, float, float], tuple[float, float, float]]: """Resolve initial camera pose from config or USD camera path.""" - if self.cfg.camera_source == "usd_path": - pose = self._resolve_camera_pose_from_usd_path(self.cfg.camera_usd_path) + if self.cfg.cam_source == "prim_path": + pose = self._resolve_camera_pose_from_usd_path(self.cfg.cam_prim_path) if pose is not None: return pose - logger.warning( - "[ViserVisualizer] camera_usd_path '%s' not found; using configured camera.", - self.cfg.camera_usd_path, + raise RuntimeError( + "[ViserVisualizer] cam_source='prim_path' requires a resolvable camera prim path, " + f"but no camera pose was found for '{self.cfg.cam_prim_path}'." ) - return self.cfg.camera_position, self.cfg.camera_target + return self._resolve_cfg_camera_pose("ViserVisualizer") def _try_apply_viser_camera_view(self, pose: tuple[tuple[float, float, float], tuple[float, float, float]]) -> bool: """Try applying camera pose to active viser clients. @@ -341,7 +344,7 @@ def _apply_pending_camera_pose(self) -> None: def _update_camera_from_usd_path(self) -> None: """Refresh camera pose from configured USD camera path when it changes.""" - pose = self._resolve_camera_pose_from_usd_path(self.cfg.camera_usd_path) + pose = self._resolve_camera_pose_from_usd_path(self.cfg.cam_prim_path) if pose is None: return if self._last_camera_pose == pose or self._pending_camera_pose == pose: diff --git a/source/isaaclab_visualizers/isaaclab_visualizers/viser/viser_visualizer_cfg.py b/source/isaaclab_visualizers/isaaclab_visualizers/viser/viser_visualizer_cfg.py index c2400c7ee1e6..f3f2aa39b0c2 100644 --- a/source/isaaclab_visualizers/isaaclab_visualizers/viser/viser_visualizer_cfg.py +++ b/source/isaaclab_visualizers/isaaclab_visualizers/viser/viser_visualizer_cfg.py @@ -35,9 +35,3 @@ class ViserVisualizerCfg(VisualizerCfg): record_to_viser: str | None = None """Path to save a .viser recording file. None = no recording.""" - - max_worlds: int | None = None - """Maximum number of worlds/environments rendered by the viewer. - - Set to ``None`` to leave this option disabled. - """ diff --git a/source/isaaclab_visualizers/setup.py b/source/isaaclab_visualizers/setup.py index 2dfe9abd30fa..fc120619787b 100644 --- a/source/isaaclab_visualizers/setup.py +++ b/source/isaaclab_visualizers/setup.py @@ -17,16 +17,16 @@ "kit": [], "newton": [ "warp-lang", - "newton @ git+https://github.com/newton-physics/newton.git@2684d75bfa4bb8b058a93b81c458a74b7701c997", + "newton @ git+https://github.com/newton-physics/newton.git@a27277ed49d6f307b8a1e4c394be7e1d14965a62", "PyOpenGL-accelerate", "imgui-bundle>=1.92.5", ], "rerun": [ - "newton @ git+https://github.com/newton-physics/newton.git@2684d75bfa4bb8b058a93b81c458a74b7701c997", + "newton @ git+https://github.com/newton-physics/newton.git@a27277ed49d6f307b8a1e4c394be7e1d14965a62", "rerun-sdk>=0.29.0", ], "viser": [ - "newton @ git+https://github.com/newton-physics/newton.git@2684d75bfa4bb8b058a93b81c458a74b7701c997", + "newton @ git+https://github.com/newton-physics/newton.git@a27277ed49d6f307b8a1e4c394be7e1d14965a62", "viser>=1.0.16", ], } diff --git a/source/isaaclab_visualizers/test/test_newton_adapter.py b/source/isaaclab_visualizers/test/test_newton_adapter.py new file mode 100644 index 000000000000..3c020a8d10ee --- /dev/null +++ b/source/isaaclab_visualizers/test/test_newton_adapter.py @@ -0,0 +1,49 @@ +# 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 + +"""Unit tests for viewer env resolution helpers.""" + +from __future__ import annotations + +from isaaclab_visualizers.newton_adapter import apply_viewer_visible_worlds, resolve_visible_env_indices + + +def test_resolve_visible_env_indices_truncates_explicit_list(): + assert resolve_visible_env_indices([1, 3, 5], 2, 10) == [1, 3] + assert resolve_visible_env_indices([1, 3], 1, 10) == [1] + + +def test_resolve_visible_env_indices_explicit_full_list_when_no_cap(): + assert resolve_visible_env_indices([1, 3], None, 10) == [1, 3] + + +def test_resolve_visible_env_indices_cap_when_no_filter(): + # When _compute_visualized_env_ids is None, cap is max_visible_envs. + assert resolve_visible_env_indices(None, 3, 10) == [0, 1, 2] + + +def test_resolve_visible_env_indices_all_when_no_cap(): + assert resolve_visible_env_indices(None, None, 10) is None + + +def test_resolve_visible_env_indices_num_envs_zero_falls_through_like_newton(): + assert resolve_visible_env_indices(None, 5, 0) is None + + +def test_apply_viewer_visible_worlds_delegates_to_resolved(): + calls: list = [] + + class _V: + def set_visible_worlds(self, worlds): + calls.append(worlds) + + apply_viewer_visible_worlds(_V(), env_ids=None, max_visible_envs=2, num_envs=5) + assert calls == [[0, 1]] + + apply_viewer_visible_worlds(_V(), env_ids=[2], max_visible_envs=99, num_envs=5) + assert calls[-1] == [2] + + apply_viewer_visible_worlds(_V(), env_ids=None, max_visible_envs=None, num_envs=3) + assert calls[-1] is None diff --git a/source/isaaclab_visualizers/test/test_visualizer_cartpole_integration.py b/source/isaaclab_visualizers/test/test_visualizer_cartpole_integration.py new file mode 100644 index 000000000000..eb5f7149432e --- /dev/null +++ b/source/isaaclab_visualizers/test/test_visualizer_cartpole_integration.py @@ -0,0 +1,608 @@ +# 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 + +"""Integration tests: cartpole env + per-backend visualizers (Kit Replicator, tiled camera, GL, Rerun, Viser). + +Visualizer packages use ``logging.getLogger(__name__)``, so loggers are named like +``isaaclab_visualizers.kit.kit_visualizer`` and ``isaaclab.visualizers.base_visualizer``. +:class:`~isaaclab.sim.simulation_context.SimulationContext` uses +``logging.getLogger(__name__)`` → ``isaaclab.sim.simulation_context``. + +We filter :class:`~pytest.LogCaptureFixture` records with :data:`_VIS_LOGGER_PREFIXES` +so only those namespaces count (not Omniverse, PhysX, or unrelated warnings). + +Set :data:`ASSERT_VISUALIZER_WARNINGS` to ``True`` locally or in CI if you want tests to +fail on WARNING-level records from those loggers; by default only ERROR+ fails. +""" + +from __future__ import annotations + +# Pyglet must use HeadlessWindow (EGL) before ``pyglet.window`` is imported so Newton +# ViewerGL can construct without an X11 display (matches ``headless=True`` on NewtonVisualizerCfg). +import pyglet + +pyglet.options["headless"] = True + +from isaaclab.app import AppLauncher + +# launch Kit app +simulation_app = AppLauncher(headless=True, enable_cameras=True).app + +import contextlib +import copy +import logging +import socket + +import numpy as np +import pytest +import torch +import warp as wp +from isaaclab_visualizers.kit import KitVisualizer, KitVisualizerCfg +from isaaclab_visualizers.newton import NewtonVisualizer, NewtonVisualizerCfg +from isaaclab_visualizers.rerun import RerunVisualizer, RerunVisualizerCfg +from isaaclab_visualizers.viser import ViserVisualizer, ViserVisualizerCfg + +import isaaclab.sim as sim_utils +from isaaclab.sim import SimulationContext + +from isaaclab_tasks.direct.cartpole.cartpole_camera_env import CartpoleCameraEnv +from isaaclab_tasks.direct.cartpole.cartpole_camera_presets_env_cfg import CartpoleCameraPresetsEnvCfg +from isaaclab_tasks.manager_based.classic.cartpole.cartpole_env_cfg import CartpolePhysicsCfg + +# When True, tests also fail on WARNING-level records from visualizer-related loggers. +ASSERT_VISUALIZER_WARNINGS = False + +_MAX_NON_BLACK_STEPS = 8 +"""Steps for tiled camera / Rerun / Viser smoke tests (early exit ok when non-black).""" + +_CARTPOLE_INTEGRATION_NUM_ENVS = 1 +"""Vectorized env count for cartpole + visualizer integration tests.""" + +_CARTPOLE_INTEGRATION_VISUALIZER_EYE: tuple[float, float, float] = (3.0, 3.0, 3.0) +"""Passed to :class:`~isaaclab.visualizers.visualizer_cfg.VisualizerCfg` subclasses (``eye``).""" + +_CARTPOLE_INTEGRATION_VISUALIZER_LOOKAT: tuple[float, float, float] = (-4.0, -4.0, 0.0) +"""Passed to visualizer cfgs (``lookat``); also applied to :class:`~isaaclab.envs.common.ViewerCfg` for the env.""" + +# Resolution overrides for this test module (cartpole preset defaults: tiled camera 100×100; Kit helper was 320×240). +_CARTPOLE_KIT_INTEGRATION_RENDER_RESOLUTION: tuple[int, int] = (600, 600) +"""Kit: Replicator ``render_product`` (width, height) for viewport RGB in the motion check.""" + +_CARTPOLE_NEWTON_INTEGRATION_WINDOW_SIZE: tuple[int, int] = (600, 600) +"""Newton: ``NewtonVisualizerCfg`` framebuffer (window_width × window_height) for ``get_frame()``.""" + +_CARTPOLE_TILED_CAMERA_INTEGRATION_WH: tuple[int, int] = (600, 600) +"""Tiled camera per-env tile width/height (preset default is 100×100); keeps ``observation_space`` consistent.""" + +_VIS_FRAME_TEST_STEPS = 60 +"""Steps for Kit / Newton frame capture: no early exit.""" + +# Motion check compares the 2nd vs last captured frame (e.g. 2nd vs 60th when *_STEPS* is 60). +_MOTION_FRAME_EARLY_IDX = 1 +"""0-based index of the *early* frame (2nd capture).""" + +_MOTION_FRAME_LATE_IDX = _VIS_FRAME_TEST_STEPS - 1 +"""0-based index of the *late* frame (e.g. 60th capture when :data:`_VIS_FRAME_TEST_STEPS` is 60).""" + +# Early vs late frame motion: void background stays similar; only count *strongly* differing pixels. +_FRAME_MOTION_CHANNEL_DIFF_THRESHOLD = 50 +"""A pixel counts as differing if max(|ΔR|, |ΔG|, |ΔB|) >= this (0–255 space).""" + +_FRAME_MOTION_MIN_DIFFERING_PIXELS = 100 +"""Minimum number of such pixels between early and late frames (stale/frozen viz should be near zero).""" + +_VIS_LOGGER_PREFIXES = ( + "isaaclab.visualizers", + "isaaclab_visualizers", + "isaaclab.sim.simulation_context", +) + + +def _logger_name_matches_visualizer_scope(logger_name: str) -> bool: + """Return True if *logger_name* is a visualizer / SimulationContext visualizer path.""" + return any(logger_name.startswith(prefix) for prefix in _VIS_LOGGER_PREFIXES) + + +def _assert_no_visualizer_log_issues(caplog: pytest.LogCaptureFixture, *, fail_on_warnings: bool | None = None) -> None: + """Fail if captured records include ERROR/CRITICAL (always) or WARNING (if *fail_on_warnings*). + + *fail_on_warnings* defaults to :data:`ASSERT_VISUALIZER_WARNINGS`. + """ + if fail_on_warnings is None: + fail_on_warnings = ASSERT_VISUALIZER_WARNINGS + + error_logs = [ + r for r in caplog.records if r.levelno >= logging.ERROR and _logger_name_matches_visualizer_scope(r.name) + ] + assert not error_logs, "Visualizer-related error logs: " + "; ".join( + f"{r.name}: {r.getMessage()}" for r in error_logs + ) + + if fail_on_warnings: + warning_logs = [ + r for r in caplog.records if r.levelno == logging.WARNING and _logger_name_matches_visualizer_scope(r.name) + ] + assert not warning_logs, "Visualizer-related warning logs: " + "; ".join( + f"{r.name}: {r.getMessage()}" for r in warning_logs + ) + + +def _configure_sim_for_visualizer_test(env: CartpoleCameraEnv) -> None: + """Settings used by the previous smoke tests; keep RTX sensors enabled for camera paths.""" + env.sim.set_setting("/isaaclab/render/rtx_sensors", True) + env.sim._app_control_on_stop_handle = None # type: ignore[attr-defined] + + +def _find_free_tcp_port(host: str = "127.0.0.1") -> int: + """Ask OS for a currently free local TCP port.""" + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.bind((host, 0)) + return int(sock.getsockname()[1]) + + +def _allocate_rerun_test_ports(host: str = "127.0.0.1") -> tuple[int, int]: + """Allocate distinct free ports for rerun web and gRPC endpoints.""" + grpc_port = _find_free_tcp_port(host) + web_port = _find_free_tcp_port(host) + while web_port == grpc_port: + web_port = _find_free_tcp_port(host) + return web_port, grpc_port + + +def _cartpole_integration_visualizer_camera_kwargs() -> dict[str, tuple[float, float, float]]: + """Eye/lookat for all :class:`~isaaclab.visualizers.visualizer_cfg.VisualizerCfg` subclasses in these tests.""" + return { + "eye": _CARTPOLE_INTEGRATION_VISUALIZER_EYE, + "lookat": _CARTPOLE_INTEGRATION_VISUALIZER_LOOKAT, + } + + +def _get_visualizer_cfg(visualizer_kind: str): + """Return (visualizer_cfg, expected_visualizer_cls) for the given visualizer kind.""" + cam = _cartpole_integration_visualizer_camera_kwargs() + if visualizer_kind == "newton": + __import__("newton") + nw, nh = _CARTPOLE_NEWTON_INTEGRATION_WINDOW_SIZE + return ( + NewtonVisualizerCfg( + headless=True, + window_width=nw, + window_height=nh, + randomly_sample_visible_envs=False, + **cam, + ), + NewtonVisualizer, + ) + if visualizer_kind == "viser": + __import__("newton") + __import__("viser") + port = _find_free_tcp_port(host="127.0.0.1") + return ( + ViserVisualizerCfg(open_browser=False, port=port, randomly_sample_visible_envs=False, **cam), + ViserVisualizer, + ) + if visualizer_kind == "rerun": + __import__("newton") + __import__("rerun") + web_port, grpc_port = _allocate_rerun_test_ports(host="127.0.0.1") + return ( + RerunVisualizerCfg( + bind_address="127.0.0.1", + open_browser=False, + web_port=web_port, + grpc_port=grpc_port, + randomly_sample_visible_envs=False, + **cam, + ), + RerunVisualizer, + ) + return KitVisualizerCfg(randomly_sample_visible_envs=False, **cam), KitVisualizer + + +def _get_physics_cfg(backend_kind: str): + """Return physics config and expected backend substring for the given backend kind.""" + if backend_kind == "physx": + __import__("isaaclab_physx") + preset = CartpolePhysicsCfg() + physics_cfg = getattr(preset, "physx", None) + if physics_cfg is None: + from isaaclab_physx.physics import PhysxCfg + + physics_cfg = PhysxCfg() + return physics_cfg, "physx" + if backend_kind == "newton": + __import__("newton") + __import__("isaaclab_newton") + preset = CartpolePhysicsCfg() + physics_cfg = getattr(preset, "newton", None) + if physics_cfg is None: + from isaaclab_newton.physics import MJWarpSolverCfg, NewtonCfg + + physics_cfg = NewtonCfg( + solver_cfg=MJWarpSolverCfg( + njmax=5, + nconmax=3, + cone="pyramidal", + impratio=1, + integrator="implicitfast", + ), + num_substeps=1, + debug_mode=False, + use_cuda_graph=True, + ) + return physics_cfg, "newton" + raise ValueError(f"Unknown backend: {backend_kind!r}") + + +def _assert_non_black_tensor(image_tensor: torch.Tensor, *, min_nonzero_pixels: int = 1) -> None: + """Assert camera-like tensor contains non-black pixels.""" + assert isinstance(image_tensor, torch.Tensor), f"Expected torch.Tensor, got {type(image_tensor)!r}" + assert image_tensor.numel() > 0, "Image tensor is empty." + finite_tensor = torch.where(torch.isfinite(image_tensor), image_tensor, torch.zeros_like(image_tensor)) + if finite_tensor.dtype.is_floating_point: + nonzero = torch.count_nonzero(torch.abs(finite_tensor) > 1e-6).item() + else: + nonzero = torch.count_nonzero(finite_tensor > 0).item() + assert nonzero >= min_nonzero_pixels, "Rendered frame appears black (no non-zero pixels)." + + +def _frame_to_numpy(frame) -> np.ndarray: + """Convert viewer ``get_frame()`` output (numpy, torch, or Warp array) to host ``numpy.ndarray``. + + ``np.asarray(wp.array)`` is unsafe: NumPy can trigger Warp indexing that raises at dimension edges. + """ + if isinstance(frame, np.ndarray): + return frame + if isinstance(frame, torch.Tensor): + return frame.detach().cpu().numpy() + if isinstance(frame, wp.array): + return wp.to_torch(frame).detach().cpu().numpy() + return np.asarray(frame) + + +def _assert_non_black_frame_array(frame) -> None: + """Assert viewer-captured frame has visible, non-black content.""" + frame_arr = _frame_to_numpy(frame) + assert frame_arr.size > 0, "Viewer returned an empty frame." + if frame_arr.ndim == 2: + color = frame_arr + else: + assert frame_arr.shape[-1] >= 3, f"Expected at least 3 channels, got shape {frame_arr.shape}." + color = frame_arr[..., :3] + finite = np.where(np.isfinite(color), color, 0) + assert np.count_nonzero(finite) > 0, "Viewer frame appears fully black." + + +def _frame_rgb_255_space(frame) -> np.ndarray: + """Return HxWx3 float in ~0–255 space for per-channel differencing.""" + arr = _frame_to_numpy(frame) + if arr.ndim == 2: + rgb = np.stack([arr, arr, arr], axis=-1) + else: + rgb = arr[..., :3] + rgb = np.asarray(rgb, dtype=np.float64) + # Normalized HDR buffers: scale so threshold matches (0,255) semantics. + if rgb.size > 0 and float(np.nanmax(rgb)) <= 1.0 + 1e-6: + rgb = rgb * 255.0 + return rgb + + +def _count_significantly_differing_pixels( + frame_a, + frame_b, + *, + channel_diff_threshold: float = _FRAME_MOTION_CHANNEL_DIFF_THRESHOLD, +) -> int: + """Count pixels where max(|ΔR|, |ΔG|, |ΔB|) >= *channel_diff_threshold* (0–255 space).""" + a = _frame_rgb_255_space(frame_a) + b = _frame_rgb_255_space(frame_b) + assert a.shape == b.shape, f"Frame shape mismatch for motion check: {a.shape} vs {b.shape}." + per_pixel_max = np.max(np.abs(a - b), axis=-1) + return int(np.count_nonzero(per_pixel_max >= channel_diff_threshold)) + + +def _assert_early_and_late_motion_frames_differ( + frames: list, + *, + channel_diff_threshold: float = _FRAME_MOTION_CHANNEL_DIFF_THRESHOLD, + min_differing_pixels: int = _FRAME_MOTION_MIN_DIFFERING_PIXELS, +) -> None: + """Fail if early vs late frames lack enough strongly differing pixels (stale/frozen bodies). + + Compares :data:`_MOTION_FRAME_EARLY_IDX` vs :data:`_MOTION_FRAME_LATE_IDX` (e.g. 2nd vs 60th capture). + + Voids/background stay near-identical; we only count pixels that change by at least + *channel_diff_threshold* on some channel (0–255). + """ + assert len(frames) >= _VIS_FRAME_TEST_STEPS, ( + f"Need at least {_VIS_FRAME_TEST_STEPS} frames for motion check, got {len(frames)}." + ) + i_early = _MOTION_FRAME_EARLY_IDX + i_late = _MOTION_FRAME_LATE_IDX + early_1 = i_early + 1 + late_1 = i_late + 1 + n_diff = _count_significantly_differing_pixels( + frames[i_early], frames[i_late], channel_diff_threshold=channel_diff_threshold + ) + assert n_diff >= min_differing_pixels, ( + f"Viewport captures #{early_1} and #{late_1} have too few strongly differing pixels " + f"({n_diff} < {min_differing_pixels}; threshold per channel={channel_diff_threshold} in 0–255 space). " + "Possible frozen or stale robot visualization." + ) + + +def _step_until_non_black_camera(env, actions: torch.Tensor, *, max_steps: int = _MAX_NON_BLACK_STEPS) -> None: + """Step env until the env's tiled camera RGB tensor is non-black, bounded by *max_steps*.""" + last_rgb = None + for _ in range(max_steps): + env.step(action=actions) + rgb = env._tiled_camera.data.output.get("rgb") + if rgb is None: + rgb = env._tiled_camera.data.output[env.cfg.tiled_camera.data_types[0]] + last_rgb = rgb + try: + _assert_non_black_tensor(rgb) + return + except AssertionError: + continue + _assert_non_black_tensor(last_rgb) + + +def _run_newton_viewer_frame_motion_test( + viewer, + *, + step_hook, + physics_kind: str, + viz_kind: str = "newton", +) -> None: + """Exactly ``_VIS_FRAME_TEST_STEPS`` sim steps; last frame non-black; early vs late motion check.""" + frames: list = [] + for _ in range(_VIS_FRAME_TEST_STEPS): + step_hook() + frames.append(viewer.get_frame()) + _assert_non_black_frame_array(frames[-1]) + _assert_early_and_late_motion_frames_differ(frames) + + +def _step_env_without_frame_check(env, actions: torch.Tensor, *, max_steps: int = _MAX_NON_BLACK_STEPS) -> None: + """Step the env to exercise visualizers that do not implement ``get_frame`` (e.g. Rerun, Viser).""" + for _ in range(max_steps): + env.step(action=actions) + + +def _build_rgb_annotator_for_camera( + camera_path: str, + *, + resolution: tuple[int, int] | None = None, +): + """Create CPU RGB annotator attached to a camera render product.""" + import omni.replicator.core as rep + + if resolution is None: + resolution = _CARTPOLE_KIT_INTEGRATION_RENDER_RESOLUTION + render_product = rep.create.render_product(camera_path, resolution=resolution) + annotator = rep.AnnotatorRegistry.get_annotator("rgb", device="cpu") + annotator.attach([render_product]) + return annotator, render_product + + +def _annotator_rgb_to_numpy(rgb_data) -> np.ndarray: + """Convert replicator annotator output to HxWx3 uint8 numpy array.""" + rgb_array = np.frombuffer(rgb_data, dtype=np.uint8).reshape(*rgb_data.shape) + if rgb_array.size == 0: + return np.zeros((1, 1, 3), dtype=np.uint8) + return rgb_array[:, :, :3] + + +def _run_kit_viewport_frame_motion_test( + env, + kit_visualizer: KitVisualizer, + *, + physics_kind: str, + viz_kind: str = "kit", +) -> None: + """Exactly ``_VIS_FRAME_TEST_STEPS`` env steps; last Replicator frame non-black; early vs late motion check.""" + camera_path = getattr(kit_visualizer, "_controlled_camera_path", None) + assert camera_path, "Kit visualizer does not expose a controlled viewport camera path." + + annotator = None + render_product = None + try: + annotator, render_product = _build_rgb_annotator_for_camera(camera_path) + actions = torch.zeros((env.num_envs, env.action_space.shape[-1]), device=env.device) + frames: list = [] + for _ in range(_VIS_FRAME_TEST_STEPS): + env.step(action=actions) + rgb_data = annotator.get_data() + frames.append(_annotator_rgb_to_numpy(rgb_data)) + _assert_non_black_frame_array(frames[-1]) + _assert_early_and_late_motion_frames_differ(frames) + finally: + if annotator is not None and render_product is not None: + with contextlib.suppress(Exception): + annotator.detach([render_product]) + + +def _make_cartpole_camera_env(visualizer_kind: str, backend_kind: str) -> CartpoleCameraEnv: + """Create cartpole camera env configured with selected visualizer and physics backend.""" + env_cfg_root = CartpoleCameraPresetsEnvCfg() + env_cfg = getattr(env_cfg_root, "default", None) + if env_cfg is None: + env_cfg = getattr(type(env_cfg_root), "default", None) + if env_cfg is None: + raise RuntimeError( + "CartpoleCameraPresetsEnvCfg does not expose a 'default' preset config. " + f"Available attributes: {sorted(vars(env_cfg_root).keys())}" + ) + env_cfg = copy.deepcopy(env_cfg) + env_cfg.scene.num_envs = _CARTPOLE_INTEGRATION_NUM_ENVS + env_cfg.viewer.eye = _CARTPOLE_INTEGRATION_VISUALIZER_EYE + env_cfg.viewer.lookat = _CARTPOLE_INTEGRATION_VISUALIZER_LOOKAT + tw, th = _CARTPOLE_TILED_CAMERA_INTEGRATION_WH + env_cfg.tiled_camera.width = tw + env_cfg.tiled_camera.height = th + if isinstance(env_cfg.observation_space, list) and len(env_cfg.observation_space) >= 3: + env_cfg.observation_space = [th, tw, env_cfg.observation_space[2]] + env_cfg.seed = None + env_cfg.sim.physics, _ = _get_physics_cfg(backend_kind) + visualizer_cfg, _ = _get_visualizer_cfg(visualizer_kind) + env_cfg.sim.visualizer_cfgs = visualizer_cfg + return CartpoleCameraEnv(env_cfg) + + +@pytest.mark.isaacsim_ci +@pytest.mark.parametrize( + "backend_kind", + [ + # xfail: Kit visualizer + PhysX only (Newton backend uses skip below — separate CUDA issue). + pytest.param( + "physx", + marks=pytest.mark.xfail( + reason=("Kit visualizer + PhysX: TODO remove xfail when stale Fabric transforms bug in Kit is fixed"), + strict=False, + ), + ), + pytest.param( + "newton", + marks=pytest.mark.skip( + reason=( + "TODO: Kit visualizer + Newton physics + Isaac RTX tiled camera can hit CUDA illegal access " + "or bad GPU state. Repro: rl_games train Isaac-Cartpole-Camera-Presets-Direct-v0 " + "--enable_cameras presets=newton --viz kit. Re-enable when fixed." + ) + ), + ), + ], +) +def test_cartpole_kit_visualizer_replicator_viewport_rgb_motion( + backend_kind: str, caplog: pytest.LogCaptureFixture +) -> None: + """Kit + cartpole: Replicator RGB on viewport camera; last frame non-black; early vs late frame differ; logs.""" + env = None + try: + sim_utils.create_new_stage() + env = _make_cartpole_camera_env(visualizer_kind="kit", backend_kind=backend_kind) + _configure_sim_for_visualizer_test(env) + with caplog.at_level(logging.WARNING): + env.reset() + kit_visualizers = [viz for viz in env.sim.visualizers if isinstance(viz, KitVisualizer)] + assert kit_visualizers, "Expected an initialized Kit visualizer." + _run_kit_viewport_frame_motion_test(env, kit_visualizers[0], physics_kind=backend_kind) + _assert_no_visualizer_log_issues(caplog) + finally: + if env is not None: + env.close() + else: + SimulationContext.clear_instance() + + +@pytest.mark.isaacsim_ci +@pytest.mark.parametrize("backend_kind", ["physx", "newton"]) +def test_cartpole_newton_visualizer_tiled_camera_rgb_non_black( + backend_kind: str, caplog: pytest.LogCaptureFixture +) -> None: + """Newton visualizer + cartpole: env tiled-camera RGB becomes non-black within a few steps; clean logs.""" + env = None + try: + sim_utils.create_new_stage() + env = _make_cartpole_camera_env(visualizer_kind="newton", backend_kind=backend_kind) + _configure_sim_for_visualizer_test(env) + with caplog.at_level(logging.WARNING): + env.reset() + actions = torch.zeros((env.num_envs, env.action_space.shape[-1]), device=env.device) + _step_until_non_black_camera(env, actions, max_steps=_MAX_NON_BLACK_STEPS) + _assert_no_visualizer_log_issues(caplog) + finally: + if env is not None: + env.close() + else: + SimulationContext.clear_instance() + + +@pytest.mark.isaacsim_ci +@pytest.mark.parametrize("backend_kind", ["physx", "newton"]) +def test_cartpole_newton_visualizer_viewergl_rgb_motion(backend_kind: str, caplog: pytest.LogCaptureFixture) -> None: + """Newton GL (``ViewerGL.get_frame``): full motion steps, last frame non-black; early vs late differ; logs.""" + env = None + try: + sim_utils.create_new_stage() + env = _make_cartpole_camera_env(visualizer_kind="newton", backend_kind=backend_kind) + _configure_sim_for_visualizer_test(env) + with caplog.at_level(logging.WARNING): + env.reset() + actions = torch.zeros((env.num_envs, env.action_space.shape[-1]), device=env.device) + newton_visualizers = [viz for viz in env.sim.visualizers if isinstance(viz, NewtonVisualizer)] + assert newton_visualizers, "Expected an initialized Newton visualizer." + viewer = getattr(newton_visualizers[0], "_viewer", None) + assert viewer is not None, "Newton viewer was not created." + + def _step_env() -> None: + env.step(action=actions) + + _run_newton_viewer_frame_motion_test(viewer, step_hook=_step_env, physics_kind=backend_kind) + _assert_no_visualizer_log_issues(caplog) + finally: + if env is not None: + env.close() + else: + SimulationContext.clear_instance() + + +@pytest.mark.isaacsim_ci +@pytest.mark.parametrize("backend_kind", ["physx", "newton"]) +def test_cartpole_rerun_visualizer_smoke_steps_and_logs(backend_kind: str, caplog: pytest.LogCaptureFixture) -> None: + """Rerun + cartpole: visualizer and viewer initialize; env steps exercise the pipeline; clean logs. + + Rerun does not expose a per-frame RGB API like ``get_frame``, so we do not assert pixel content. + """ + env = None + try: + sim_utils.create_new_stage() + env = _make_cartpole_camera_env(visualizer_kind="rerun", backend_kind=backend_kind) + _configure_sim_for_visualizer_test(env) + with caplog.at_level(logging.WARNING): + env.reset() + actions = torch.zeros((env.num_envs, env.action_space.shape[-1]), device=env.device) + rerun_visualizers = [viz for viz in env.sim.visualizers if isinstance(viz, RerunVisualizer)] + assert rerun_visualizers, "Expected an initialized Rerun visualizer." + assert getattr(rerun_visualizers[0], "_viewer", None) is not None, "Rerun viewer was not created." + _step_env_without_frame_check(env, actions, max_steps=_MAX_NON_BLACK_STEPS) + _assert_no_visualizer_log_issues(caplog) + finally: + if env is not None: + env.close() + else: + SimulationContext.clear_instance() + + +@pytest.mark.isaacsim_ci +@pytest.mark.parametrize("backend_kind", ["physx", "newton"]) +def test_cartpole_viser_visualizer_smoke_steps_and_logs(backend_kind: str, caplog: pytest.LogCaptureFixture) -> None: + """Viser + cartpole: visualizer and viewer initialize; env steps exercise the pipeline; clean logs. + + No per-frame RGB assertion (Viser does not mirror the Newton ``get_frame`` path used elsewhere). + """ + env = None + try: + sim_utils.create_new_stage() + env = _make_cartpole_camera_env(visualizer_kind="viser", backend_kind=backend_kind) + _configure_sim_for_visualizer_test(env) + with caplog.at_level(logging.WARNING): + env.reset() + actions = torch.zeros((env.num_envs, env.action_space.shape[-1]), device=env.device) + viser_visualizers = [viz for viz in env.sim.visualizers if isinstance(viz, ViserVisualizer)] + assert viser_visualizers, "Expected an initialized Viser visualizer." + assert getattr(viser_visualizers[0], "_viewer", None) is not None, "Viser viewer was not created." + _step_env_without_frame_check(env, actions, max_steps=_MAX_NON_BLACK_STEPS) + _assert_no_visualizer_log_issues(caplog) + finally: + if env is not None: + env.close() + else: + SimulationContext.clear_instance() + + +if __name__ == "__main__": + pytest.main([__file__, "-v", "--maxfail=1"]) diff --git a/source/isaaclab_visualizers/test/test_visualizer_smoke_logs.py b/source/isaaclab_visualizers/test/test_visualizer_smoke_logs.py deleted file mode 100644 index 22f620fb02a8..000000000000 --- a/source/isaaclab_visualizers/test/test_visualizer_smoke_logs.py +++ /dev/null @@ -1,228 +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 - -"""Smoke test visualizer stepping and error logging.""" - -from isaaclab.app import AppLauncher - -# launch Kit app -simulation_app = AppLauncher(headless=True, enable_cameras=True).app - -import logging -import socket - -import pytest -import torch -from isaaclab_visualizers.kit import KitVisualizer, KitVisualizerCfg -from isaaclab_visualizers.newton import NewtonVisualizer, NewtonVisualizerCfg -from isaaclab_visualizers.rerun import RerunVisualizer, RerunVisualizerCfg -from isaaclab_visualizers.viser import ViserVisualizer, ViserVisualizerCfg - -import isaaclab.sim as sim_utils -from isaaclab.envs import DirectRLEnv, DirectRLEnvCfg -from isaaclab.scene import InteractiveSceneCfg -from isaaclab.sim import SimulationCfg, SimulationContext -from isaaclab.utils import configclass - -from isaaclab_tasks.manager_based.classic.cartpole.cartpole_env_cfg import ( - CartpolePhysicsCfg, - CartpoleSceneCfg, -) - -# Set to False to only fail on visualizer errors; when True, also fail on warnings. -ASSERT_VISUALIZER_WARNINGS = True - -_SMOKE_STEPS = 4 -_VIS_LOGGER_PREFIXES = ( - "isaaclab.visualizers", - "isaaclab_visualizers", - "isaaclab.sim.simulation_context", -) - - -def _find_free_tcp_port(host: str = "127.0.0.1") -> int: - """Ask OS for a currently free local TCP port.""" - with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: - sock.bind((host, 0)) - return int(sock.getsockname()[1]) - - -def _allocate_rerun_test_ports(host: str = "127.0.0.1") -> tuple[int, int]: - """Allocate distinct free ports for rerun web and gRPC endpoints.""" - grpc_port = _find_free_tcp_port(host) - web_port = _find_free_tcp_port(host) - while web_port == grpc_port: - web_port = _find_free_tcp_port(host) - return web_port, grpc_port - - -@configclass -class _SmokeEnvCfg(DirectRLEnvCfg): - decimation: int = 2 - action_space: int = 0 - observation_space: int = 0 - episode_length_s: float = 5.0 - sim: SimulationCfg = SimulationCfg(dt=0.005, render_interval=2, visualizer_cfgs=KitVisualizerCfg()) - scene: InteractiveSceneCfg = InteractiveSceneCfg(num_envs=1, env_spacing=1.0) - - -class _SmokeEnv(DirectRLEnv): - def _pre_physics_step(self, actions): - return - - def _apply_action(self): - return - - def _get_observations(self): - return {} - - def _get_rewards(self): - return {} - - def _get_dones(self): - return torch.zeros(1, dtype=torch.bool), torch.zeros(1, dtype=torch.bool) - - -def _get_visualizer_cfg(visualizer_kind: str): - """Return (visualizer_cfg, expected_visualizer_cls) for the given visualizer kind.""" - if visualizer_kind == "newton": - __import__("newton") - return NewtonVisualizerCfg(headless=True), NewtonVisualizer - if visualizer_kind == "viser": - __import__("newton") - __import__("viser") - return ViserVisualizerCfg(open_browser=False), ViserVisualizer - if visualizer_kind == "rerun": - __import__("newton") - __import__("rerun") - web_port, grpc_port = _allocate_rerun_test_ports(host="127.0.0.1") - # Use dynamically allocated non-default ports in smoke tests to avoid collisions. - # TODO: Consider supporting cleanup/termination of stale rerun processes when ports are occupied. - return ( - RerunVisualizerCfg( - bind_address="127.0.0.1", - open_browser=False, - web_port=web_port, - grpc_port=grpc_port, - ), - RerunVisualizer, - ) - return KitVisualizerCfg(), KitVisualizer - - -def _get_physics_cfg(backend_kind: str): - """Return physics config and expected backend substring for the given backend kind. - - Uses cartpole preset instance so we work whether presets are class or instance attributes. - Fallback: build PhysxCfg/NewtonCfg in-test if preset does not expose that backend. - """ - if backend_kind == "physx": - __import__("isaaclab_physx") - preset = CartpolePhysicsCfg() - physics_cfg = getattr(preset, "physx", None) - if physics_cfg is None: - from isaaclab_physx.physics import PhysxCfg - - physics_cfg = PhysxCfg() - return physics_cfg, "physx" - if backend_kind == "newton": - __import__("newton") - __import__("isaaclab_newton") - preset = CartpolePhysicsCfg() - physics_cfg = getattr(preset, "newton", None) - if physics_cfg is None: - from isaaclab_newton.physics import MJWarpSolverCfg, NewtonCfg - - physics_cfg = NewtonCfg( - solver_cfg=MJWarpSolverCfg( - njmax=5, - nconmax=3, - cone="pyramidal", - impratio=1, - integrator="implicitfast", - ), - num_substeps=1, - debug_mode=False, - use_cuda_graph=True, - ) - return physics_cfg, "newton" - raise ValueError(f"Unknown backend: {backend_kind!r}") - - -def _resolve_case(visualizer_kind: str, backend_kind: str): - """Resolve (env_cfg, expected_visualizer_cls, expected_backend_substring) for one smoke test. - - Uses cartpole scene for all combinations (works with both PhysX and Newton). - """ - scene_cfg = CartpoleSceneCfg(num_envs=1, env_spacing=1.0) - viz_cfg, expected_viz_cls = _get_visualizer_cfg(visualizer_kind) - physics_cfg, expected_backend = _get_physics_cfg(backend_kind) - - cfg = _SmokeEnvCfg() - cfg.scene = scene_cfg - cfg.sim = SimulationCfg( - dt=0.005, - render_interval=2, - visualizer_cfgs=viz_cfg, - physics=physics_cfg, - ) - return cfg, expected_viz_cls, expected_backend - - -def _run_smoke_test(cfg, expected_visualizer_cls, expected_backend: str, caplog) -> None: - """Run smoke steps and assert no visualizer errors; optionally no warnings (see ASSERT_VISUALIZER_WARNINGS).""" - env = None - try: - sim_utils.create_new_stage() - env = _SmokeEnv(cfg=cfg) - backend_name = env.sim.physics_manager.__name__.lower() - assert expected_backend in backend_name, ( - f"Expected physics backend containing {expected_backend!r}, got {backend_name!r}" - ) - env.sim.set_setting("/isaaclab/render/rtx_sensors", True) - env.sim._app_control_on_stop_handle = None # type: ignore[attr-defined] - - actions = torch.zeros((env.num_envs, 0), device=env.device) - with caplog.at_level(logging.WARNING): - env.reset() - assert env.sim.visualizers - assert isinstance(env.sim.visualizers[0], expected_visualizer_cls) - for _ in range(_SMOKE_STEPS): - env.step(action=actions) - - # Always fail on errors - error_logs = [ - r for r in caplog.records if r.levelno >= logging.ERROR and r.name.startswith(_VIS_LOGGER_PREFIXES) - ] - assert not error_logs, "Visualizer emitted error logs during smoke stepping: " + "; ".join( - f"{r.name}: {r.message}" for r in error_logs - ) - - # Optionally fail on warnings - if ASSERT_VISUALIZER_WARNINGS: - warning_logs = [ - r for r in caplog.records if r.levelno >= logging.WARNING and r.name.startswith(_VIS_LOGGER_PREFIXES) - ] - assert not warning_logs, "Visualizer emitted warning logs during smoke stepping: " + "; ".join( - f"{r.name}: {r.message}" for r in warning_logs - ) - finally: - if env is not None: - env.close() - else: - SimulationContext.clear_instance() - - -@pytest.mark.isaacsim_ci -@pytest.mark.parametrize("visualizer_kind", ["kit", "newton", "rerun", "viser"]) -@pytest.mark.parametrize("backend_kind", ["physx", "newton"]) -def test_visualizer_backend_smoke(visualizer_kind: str, backend_kind: str, caplog): - """Smoke test each (visualizer, backend) pair; assert no errors (optionally no warnings).""" - cfg, expected_viz_cls, expected_backend = _resolve_case(visualizer_kind, backend_kind) - _run_smoke_test(cfg, expected_viz_cls, expected_backend, caplog) - - -if __name__ == "__main__": - pytest.main([__file__, "-v", "--maxfail=1"]) diff --git a/tools/wheel_builder/res/python_packages.toml b/tools/wheel_builder/res/python_packages.toml index d79ce41ada84..f6a42b90a1bc 100644 --- a/tools/wheel_builder/res/python_packages.toml +++ b/tools/wheel_builder/res/python_packages.toml @@ -83,9 +83,9 @@ pyproject.optional-dependencies.all = [ # ================================================================================ { "newton" = [ "warp-lang==1.12.0", - "mujoco==3.5.0", - "mujoco-warp==3.5.0.2", - "newton==1.0.0", + "mujoco==3.6.0", + "mujoco-warp==3.6.0", + "newton @ git+https://github.com/newton-physics/newton.git@a27277ed49d6f307b8a1e4c394be7e1d14965a62", "PyOpenGL-accelerate==3.1.10" ] }, # ================================================================================