diff --git a/docs/source/testing/index.rst b/docs/source/testing/index.rst new file mode 100644 index 00000000000..e1d00fa97ff --- /dev/null +++ b/docs/source/testing/index.rst @@ -0,0 +1,12 @@ +.. _testing: + +Testing +======= + +This section covers testing utilities and patterns for Isaac Lab development. + +.. toctree:: + :maxdepth: 2 + + mock_interfaces + micro_benchmarks diff --git a/docs/source/testing/micro_benchmarks.rst b/docs/source/testing/micro_benchmarks.rst new file mode 100644 index 00000000000..9f2dc1fb4d9 --- /dev/null +++ b/docs/source/testing/micro_benchmarks.rst @@ -0,0 +1,348 @@ +.. _testing_micro_benchmarks: + +Micro-Benchmarks for Performance Testing +======================================== + +Isaac Lab provides micro-benchmarking tools to measure the performance of asset +setter/writer methods and data property accessors without requiring Isaac Sim. + +Overview +-------- + +The benchmarks use **mock interfaces** to simulate PhysX views, allowing performance +measurement of Python-level overhead in isolation. This is useful for: + +- Comparing list vs tensor index performance +- Identifying bottlenecks in hot code paths +- Tracking performance regressions +- Optimizing custom methods + +Quick Start +----------- + +Run benchmarks using the Isaac Lab launcher: + +.. code-block:: bash + + # Run Articulation method benchmarks + ./isaaclab.sh -p source/isaaclab_physx/benchmark/assets/benchmark_articulation.py + + # With custom parameters + ./isaaclab.sh -p source/isaaclab_physx/benchmark/assets/benchmark_articulation.py \ + --num_iterations 1000 \ + --num_instances 64 \ + --num_bodies 5 \ + --num_joints 4 + +Available Benchmarks +-------------------- + +Asset Method Benchmarks +~~~~~~~~~~~~~~~~~~~~~~~ + +These benchmark setter and writer methods on asset classes: + +.. list-table:: + :header-rows: 1 + :widths: 35 25 40 + + * - Benchmark File + - Asset Class + - Methods Covered + * - ``benchmark_articulation.py`` + - ``Articulation`` + - 24 methods (root/joint state, mass props, forces) + * - ``benchmark_rigid_object.py`` + - ``RigidObject`` + - 13 methods (root state, mass props, forces) + * - ``benchmark_rigid_object_collection.py`` + - ``RigidObjectCollection`` + - 13 methods (body state, mass props, forces) + +Data Property Benchmarks +~~~~~~~~~~~~~~~~~~~~~~~~ + +These benchmark property accessors on data classes: + +.. list-table:: + :header-rows: 1 + :widths: 40 30 30 + + * - Benchmark File + - Data Class + - Properties + * - ``benchmark_articulation_data.py`` + - ``ArticulationData`` + - 59 properties + * - ``benchmark_rigid_object_data.py`` + - ``RigidObjectData`` + - 40 properties + * - ``benchmark_rigid_object_collection_data.py`` + - ``RigidObjectCollectionData`` + - 40 properties + +All benchmarks are located in ``source/isaaclab_physx/benchmark/assets/``. + +Command Line Arguments +---------------------- + +Common Arguments +~~~~~~~~~~~~~~~~ + +.. list-table:: + :header-rows: 1 + :widths: 20 15 65 + + * - Argument + - Default + - Description + * - ``--num_iterations`` + - 1000 + - Number of timed iterations + * - ``--warmup_steps`` + - 10 + - Warmup iterations (not timed) + * - ``--num_instances`` + - 4096 + - Number of asset instances + * - ``--device`` + - ``cuda:0`` + - Device for tensors + * - ``--mode`` + - ``all`` + - ``all``, ``torch_list``, or ``torch_tensor`` + * - ``--output`` + - auto + - Output JSON filename + * - ``--no_csv`` + - false + - Disable CSV output + +Asset-Specific Arguments +~~~~~~~~~~~~~~~~~~~~~~~~ + +**Articulation benchmarks:** + +- ``--num_bodies``: Number of links (default: 13) +- ``--num_joints``: Number of DOFs (default: 12) + +**RigidObjectCollection benchmarks:** + +- ``--num_bodies``: Number of bodies in collection (default: 5) + +Benchmark Modes +--------------- + +Each method is benchmarked under two input scenarios: + +**torch_list** + Environment/body IDs passed as Python lists. Measures the overhead of + list-to-tensor conversion, which is common in user code. + +**torch_tensor** + Environment/body IDs passed as pre-allocated tensors. Represents the + optimal baseline with minimal overhead. + +Example output: + +.. code-block:: text + + [1/24] [TORCH_LIST] write_root_state_to_sim... 132.02 ± 6.79 µs + [1/24] [TORCH_TENSOR] write_root_state_to_sim... 65.44 ± 3.06 µs + +The comparison shows tensor indices are ~2x faster than list indices. + +Output Format +------------- + +Console Output +~~~~~~~~~~~~~~ + +.. code-block:: text + + Benchmarking Articulation (PhysX) with 64 instances, 5 bodies, 4 joints... + Device: cuda:0 + Iterations: 100, Warmup: 10 + + Benchmarking 24 methods... + [1/24] [TORCH_LIST] write_root_state_to_sim... 132.02 ± 6.79 µs + [1/24] [TORCH_TENSOR] write_root_state_to_sim... 65.44 ± 3.06 µs + ... + + ================================================================================ + COMPARISON: Torch_list vs Torch_tensor + ================================================================================ + Method Name Torch_list Torch_tensor Speedup + ------------------------------------------------------------------------ + write_root_state_to_sim 132.02 65.44 2.02x + +Export Files +~~~~~~~~~~~~ + +Results are automatically exported to: + +- ``{benchmark_name}_{timestamp}.json`` - Full results with hardware info +- ``{benchmark_name}_{timestamp}.csv`` - Tabular results for analysis + +JSON Structure +~~~~~~~~~~~~~~ + +.. code-block:: json + + { + "config": { + "num_iterations": 100, + "num_instances": 64, + "device": "cuda:0" + }, + "hardware": { + "cpu": "Intel Core i9-13950HX", + "gpu": "NVIDIA RTX 5000", + "pytorch": "2.7.0", + "cuda": "12.8" + }, + "results": [ + { + "name": "write_root_state_to_sim", + "mode": "torch_list", + "mean_us": 132.02, + "std_us": 6.79, + "iterations": 100 + } + ] + } + +Architecture +------------ + +The benchmarks use mock interfaces to simulate PhysX views without Isaac Sim: + +.. code-block:: text + + ┌─────────────────────┐ ┌──────────────────────┐ + │ Asset Class │────>│ MockArticulationView│ + │ (Articulation) │ │ (mock_interfaces) │ + └─────────────────────┘ └──────────────────────┘ + │ + v + ┌─────────────────────┐ + │ Benchmark │ + │ Framework │ + └─────────────────────┘ + +Key Components +~~~~~~~~~~~~~~ + +1. **Mock Views** (``isaaclab_physx/test/mock_interfaces/``) + + - ``MockArticulationView`` - Mimics PhysX ArticulationView + - ``MockRigidBodyView`` - Mimics PhysX RigidBodyView + +2. **Benchmark Utilities** (``isaaclab/test/benchmark/``) + + - ``BenchmarkConfig`` - Configuration dataclass + - ``MethodBenchmark`` - Benchmark definition + - ``benchmark_method()`` - Core timing function + - Export utilities for JSON/CSV + +3. **Module Mocking** + + Each benchmark file mocks Isaac Sim dependencies (``isaacsim``, ``omni``, ``pxr``) + to allow the asset classes to be instantiated without simulation. + +Adding New Benchmarks +--------------------- + +Adding a Method Benchmark +~~~~~~~~~~~~~~~~~~~~~~~~~ + +1. Create input generator functions: + +.. code-block:: python + + def gen_my_method_torch_list(config: BenchmarkConfig) -> dict: + return { + "param1": torch.rand(config.num_instances, 3, device=config.device), + "env_ids": list(range(config.num_instances)), + } + + def gen_my_method_torch_tensor(config: BenchmarkConfig) -> dict: + return { + "param1": torch.rand(config.num_instances, 3, device=config.device), + "env_ids": make_tensor_env_ids(config.num_instances, config.device), + } + +2. Add to the ``BENCHMARKS`` list: + +.. code-block:: python + + MethodBenchmark( + name="my_method", + method_name="my_method", + input_generators={ + "torch_list": gen_my_method_torch_list, + "torch_tensor": gen_my_method_torch_tensor, + }, + category="my_category", + ), + +Adding a Property Benchmark +~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +For data class properties, add to the ``PROPERTIES`` list: + +.. code-block:: python + + ("my_property", {"derived_from": ["dependency1", "dependency2"]}), + +The ``derived_from`` key indicates dependencies that should be pre-computed +before timing the property access. + +Performance Tips +---------------- + +Based on benchmark results: + +1. **Use tensor indices** instead of lists for 30-50% speedup +2. **Pre-allocate index tensors** and reuse them across calls +3. **Batch operations** where possible (e.g., set all joint positions at once) +4. **Mass properties are CPU-bound** - PhysX requires CPU tensors for these + +Example optimization: + +.. code-block:: python + + # Slow: Create new list each call + for _ in range(1000): + robot.write_joint_state_to_sim(state, env_ids=list(range(64))) + + # Fast: Pre-allocate tensor and reuse + env_ids = torch.arange(64, device="cuda:0") + for _ in range(1000): + robot.write_joint_state_to_sim(state, env_ids=env_ids) + +Troubleshooting +--------------- + +Import Errors +~~~~~~~~~~~~~ + +Ensure you're running through ``isaaclab.sh``: + +.. code-block:: bash + + ./isaaclab.sh -p source/isaaclab_physx/benchmark/assets/benchmark_articulation.py + +CUDA Out of Memory +~~~~~~~~~~~~~~~~~~ + +Reduce ``--num_instances``: + +.. code-block:: bash + + ./isaaclab.sh -p ... --num_instances 1024 + +Slow First Run +~~~~~~~~~~~~~~ + +The first run compiles Warp kernels. Subsequent runs will be faster. diff --git a/docs/source/testing/mock_interfaces.rst b/docs/source/testing/mock_interfaces.rst new file mode 100644 index 00000000000..63c9a2477db --- /dev/null +++ b/docs/source/testing/mock_interfaces.rst @@ -0,0 +1,470 @@ +.. _testing_mock_interfaces: + +Mock Interfaces for Unit Testing +================================ + +Isaac Lab provides mock implementations of sensor, asset, and PhysX view classes +for unit testing without requiring Isaac Sim or GPU simulation. + +Overview +-------- + +Two levels of mock interfaces are available: + +1. **High-level mocks** (``isaaclab.test.mock_interfaces``) - Mock sensors and assets +2. **Low-level PhysX mocks** (``isaaclab_physx.test.mock_interfaces``) - Mock PhysX TensorAPI views + +High-Level Mock Interfaces +-------------------------- + +Located in ``isaaclab.test.mock_interfaces``, these mock the public API of sensor +and asset base classes. + +Available Mocks +~~~~~~~~~~~~~~~ + +**Sensors:** + +.. list-table:: + :header-rows: 1 + :widths: 25 25 50 + + * - Mock Class + - Real Class + - Description + * - ``MockImu`` + - ``BaseImu`` + - IMU sensor (accelerometer + gyroscope) + * - ``MockContactSensor`` + - ``BaseContactSensor`` + - Contact force sensor + * - ``MockFrameTransformer`` + - ``BaseFrameTransformer`` + - Frame tracking sensor + +**Assets:** + +.. list-table:: + :header-rows: 1 + :widths: 25 25 50 + + * - Mock Class + - Real Class + - Description + * - ``MockArticulation`` + - ``BaseArticulation`` + - Articulated robot (joints + bodies) + * - ``MockRigidObject`` + - ``BaseRigidObject`` + - Single rigid body + * - ``MockRigidObjectCollection`` + - ``BaseRigidObjectCollection`` + - Collection of rigid bodies + +Quick Start +~~~~~~~~~~~ + +.. code-block:: python + + import torch + from isaaclab.test.mock_interfaces import MockArticulation, MockContactSensor + + # Create a mock quadruped robot + robot = MockArticulation( + num_instances=4, + num_joints=12, + num_bodies=13, + joint_names=[f"joint_{i}" for i in range(12)], + device="cpu" + ) + + # Set joint states + robot.data.set_mock_data( + joint_pos=torch.zeros(4, 12), + joint_vel=torch.randn(4, 12) * 0.1 + ) + + # Access data like the real asset + positions = robot.data.joint_pos # Shape: (4, 12) + + # Create a mock contact sensor + sensor = MockContactSensor(num_instances=4, num_bodies=4, device="cpu") + sensor.data.set_net_forces_w(torch.randn(4, 4, 3)) + +Factory Functions +~~~~~~~~~~~~~~~~~ + +Pre-configured factories for common use cases: + +.. code-block:: python + + from isaaclab.test.mock_interfaces import ( + create_mock_quadruped, + create_mock_humanoid, + create_mock_foot_contact_sensor, + ) + + # Pre-configured quadruped (12 joints, 13 bodies) + robot = create_mock_quadruped(num_instances=4) + + # Pre-configured humanoid (21 joints, 22 bodies) + humanoid = create_mock_humanoid(num_instances=2) + + # Foot contact sensor for quadruped + foot_contact = create_mock_foot_contact_sensor(num_instances=4, num_feet=4) + +Patching Utilities +~~~~~~~~~~~~~~~~~~ + +Use context managers or decorators to inject mocks into tests: + +**Context Managers:** + +.. code-block:: python + + from isaaclab.test.mock_interfaces.utils import patch_articulation, patch_sensor + + with patch_articulation("my_module.Articulation", num_joints=12) as MockRobot: + robot = my_function_that_creates_robot() + robot.data.set_joint_pos(torch.zeros(1, 12)) + + with patch_sensor("my_module.ContactSensor", "contact", num_bodies=4): + sensor = my_function_that_creates_sensor() + +**Decorators:** + +.. code-block:: python + + from isaaclab.test.mock_interfaces.utils import mock_articulation, mock_sensor + + @mock_articulation(num_joints=12, num_bodies=13) + def test_observation_function(mock_robot): + mock_robot.data.set_joint_pos(torch.zeros(1, 12)) + obs = compute_observation(mock_robot) + assert obs.shape == (1, 24) + + @mock_sensor("contact", num_instances=4, num_bodies=4) + def test_contact_reward(mock_contact): + mock_contact.data.set_net_forces_w(torch.randn(4, 4, 3)) + reward = compute_contact_reward(mock_contact) + +Low-Level PhysX Mock Interfaces +------------------------------- + +Located in ``isaaclab_physx.test.mock_interfaces``, these mock the PhysX TensorAPI +views used internally by assets and sensors. + +Available Mocks +~~~~~~~~~~~~~~~ + +.. list-table:: + :header-rows: 1 + :widths: 25 25 50 + + * - Mock Class + - PhysX Class + - Used By + * - ``MockRigidBodyView`` + - ``physx.RigidBodyView`` + - ``RigidObject`` + * - ``MockArticulationView`` + - ``physx.ArticulationView`` + - ``Articulation`` + * - ``MockRigidContactView`` + - ``physx.RigidContactView`` + - ``ContactSensor`` + +Quick Start +~~~~~~~~~~~ + +.. code-block:: python + + from isaaclab_physx.test.mock_interfaces import ( + MockRigidBodyView, + MockArticulationView, + MockRigidContactView, + ) + + # Create a mock rigid body view + view = MockRigidBodyView(count=4, device="cpu") + transforms = view.get_transforms() # Shape: (4, 7) + + # Set mock data + view.set_mock_transforms(torch.randn(4, 7)) + + # Create a mock articulation view + art_view = MockArticulationView( + count=4, + num_dofs=12, + num_links=13, + device="cpu" + ) + positions = art_view.get_dof_positions() # Shape: (4, 12) + +Factory Functions +~~~~~~~~~~~~~~~~~ + +.. code-block:: python + + from isaaclab_physx.test.mock_interfaces import ( + create_mock_rigid_body_view, + create_mock_articulation_view, + create_mock_quadruped_view, + create_mock_rigid_contact_view, + ) + + # Pre-configured quadruped view (12 DOFs, 13 links) + view = create_mock_quadruped_view(count=4) + + # Pre-configured humanoid view (21 DOFs, 22 links) + humanoid = create_mock_humanoid_view(count=2) + +Patching Utilities +~~~~~~~~~~~~~~~~~~ + +.. code-block:: python + + from isaaclab_physx.test.mock_interfaces.utils import ( + patch_rigid_body_view, + patch_articulation_view, + mock_articulation_view, + ) + + # Context manager + with patch_rigid_body_view("my_module.physx.RigidBodyView", count=4): + view = create_rigid_body() + + # Decorator + @mock_articulation_view(count=4, num_dofs=12, num_links=13) + def test_my_function(mock_view): + positions = mock_view.get_dof_positions() + assert positions.shape == (4, 12) + +Data Shapes Reference +--------------------- + +High-Level Mocks +~~~~~~~~~~~~~~~~ + +**IMU Data:** + +.. list-table:: + :header-rows: 1 + :widths: 25 20 55 + + * - Property + - Shape + - Description + * - ``pos_w`` + - ``(N, 3)`` + - Position in world frame + * - ``quat_w`` + - ``(N, 4)`` + - Orientation (w,x,y,z) in world frame + * - ``lin_vel_b`` + - ``(N, 3)`` + - Linear velocity in body frame + * - ``ang_vel_b`` + - ``(N, 3)`` + - Angular velocity in body frame + * - ``projected_gravity_b`` + - ``(N, 3)`` + - Gravity in body frame + +**Contact Sensor Data:** + +.. list-table:: + :header-rows: 1 + :widths: 25 20 55 + + * - Property + - Shape + - Description + * - ``net_forces_w`` + - ``(N, B, 3)`` + - Net contact forces + * - ``force_matrix_w`` + - ``(N, B, M, 3)`` + - Filtered forces + * - ``current_contact_time`` + - ``(N, B)`` + - Time in contact + * - ``current_air_time`` + - ``(N, B)`` + - Time in air + +Where N = instances, B = bodies, M = filter bodies + +**Articulation Data:** + +.. list-table:: + :header-rows: 1 + :widths: 25 20 55 + + * - Property + - Shape + - Description + * - ``joint_pos`` + - ``(N, J)`` + - Joint positions + * - ``joint_vel`` + - ``(N, J)`` + - Joint velocities + * - ``root_link_pose_w`` + - ``(N, 7)`` + - Root pose in world + * - ``body_link_pose_w`` + - ``(N, B, 7)`` + - Body poses in world + +Where N = instances, J = joints, B = bodies + +PhysX View Mocks +~~~~~~~~~~~~~~~~ + +**RigidBodyView:** + +.. list-table:: + :header-rows: 1 + :widths: 30 20 50 + + * - Method + - Shape + - Description + * - ``get_transforms()`` + - ``(N, 7)`` + - [pos(3), quat_xyzw(4)] + * - ``get_velocities()`` + - ``(N, 6)`` + - [lin_vel(3), ang_vel(3)] + * - ``get_masses()`` + - ``(N, 1, 1)`` + - mass per body + * - ``get_inertias()`` + - ``(N, 1, 3, 3)`` + - 3x3 inertia matrix + +**ArticulationView:** + +.. list-table:: + :header-rows: 1 + :widths: 35 20 45 + + * - Method + - Shape + - Description + * - ``get_root_transforms()`` + - ``(N, 7)`` + - Root pose + * - ``get_link_transforms()`` + - ``(N, L, 7)`` + - Link poses + * - ``get_dof_positions()`` + - ``(N, J)`` + - Joint positions + * - ``get_dof_limits()`` + - ``(N, J, 2)`` + - [lower, upper] + +**RigidContactView:** + +.. list-table:: + :header-rows: 1 + :widths: 40 20 40 + + * - Method + - Shape + - Description + * - ``get_net_contact_forces(dt)`` + - ``(N*B, 3)`` + - Flat net forces + * - ``get_contact_force_matrix(dt)`` + - ``(N*B, F, 3)`` + - Per-filter forces + +Design Patterns +--------------- + +All mock interfaces follow these patterns: + +1. **Lazy Initialization** - Tensors created on first access with correct shapes +2. **Device Transfer** - All setters call ``.to(device)`` +3. **Identity Quaternion Defaults** - Quaternions default to identity +4. **Clone on Getters** - Return ``.clone()`` to prevent mutation +5. **No-op Actions** - Simulation operations (reset, update, write_to_sim) do nothing +6. **Mock Setters** - Direct ``set_mock_*`` methods for test setup + +Example: Testing an Observation Function +---------------------------------------- + +.. code-block:: python + + import pytest + import torch + from isaaclab.test.mock_interfaces import MockArticulation + + @pytest.fixture + def robot(): + return MockArticulation( + num_instances=4, + num_joints=12, + num_bodies=13, + device="cpu" + ) + + def test_joint_observation(robot): + # Setup + joint_pos = torch.randn(4, 12) + robot.data.set_joint_pos(joint_pos) + + # Test your observation function + obs = compute_joint_observation(robot) + + # Verify + assert obs.shape == (4, 24) # pos + vel + assert torch.allclose(obs[:, :12], joint_pos) + + def test_body_observation(robot): + # Setup + body_poses = torch.randn(4, 13, 7) + robot.data.set_body_link_pose_w(body_poses) + + # Test + obs = compute_body_observation(robot) + assert obs.shape == (4, 13, 7) + +Example: Testing with PhysX View Mocks +-------------------------------------- + +.. code-block:: python + + import pytest + import torch + from isaaclab_physx.test.mock_interfaces import create_mock_quadruped_view + from isaaclab_physx.test.mock_interfaces.utils import mock_articulation_view + + def test_quadruped_joint_limits(): + """Test quadruped joint limit handling.""" + view = create_mock_quadruped_view(count=4) + + # Set custom joint limits + limits = torch.zeros(4, 12, 2) + limits[:, :, 0] = -1.5 # lower limit + limits[:, :, 1] = 1.5 # upper limit + view.set_mock_dof_limits(limits) + + # Verify limits + result = view.get_dof_limits() + assert torch.allclose(result[:, :, 0], torch.full((4, 12), -1.5)) + assert torch.allclose(result[:, :, 1], torch.full((4, 12), 1.5)) + + @mock_articulation_view(count=4, num_dofs=12, num_links=13) + def test_articulation_with_decorator(mock_view): + """Test using the decorator pattern.""" + # Set some initial positions + positions = torch.randn(4, 12) + mock_view.set_mock_dof_positions(positions) + + # Verify + result = mock_view.get_dof_positions() + assert torch.allclose(result, positions) diff --git a/source/isaaclab/config/extension.toml b/source/isaaclab/config/extension.toml index f56b9f1a68a..2b7bb206646 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 = "2.1.0" +version = "2.1.2" # Description title = "Isaac Lab framework for Robot Learning" diff --git a/source/isaaclab/docs/CHANGELOG.rst b/source/isaaclab/docs/CHANGELOG.rst index 3b59fa50e19..1f798735dc7 100644 --- a/source/isaaclab/docs/CHANGELOG.rst +++ b/source/isaaclab/docs/CHANGELOG.rst @@ -1,6 +1,48 @@ Changelog --------- +2.1.2 (2026-01-30) +~~~~~~~~~~~~~~~~~~~ + +Added +^^^^^ + +* Added :mod:`isaaclab.test.benchmark` module providing a comprehensive benchmarking framework + for measuring performance of Isaac Lab components. Includes: + + * :class:`BenchmarkConfig`: Configuration dataclass for benchmark execution parameters + (iterations, warmup steps, instances, device). + * :class:`BenchmarkResult`: Dataclass capturing timing statistics (mean, std in microseconds), + skip status, and dependency information. + * :class:`MethodBenchmark`: Definition class for methods to benchmark with multi-mode + input generators. + * Input generator helpers for creating standardized tensors and Warp masks: + ``make_tensor_env_ids``, ``make_tensor_joint_ids``, ``make_tensor_body_ids``, + ``make_warp_env_mask``, ``make_warp_joint_mask``, ``make_warp_body_mask``. + * :func:`benchmark_method`: Core function for benchmarking with warmup phases, + GPU synchronization, and graceful error handling. + * I/O utilities: :func:`get_hardware_info`, :func:`get_git_info`, :func:`print_hardware_info`, + :func:`print_results`, :func:`export_results_json`, :func:`export_results_csv`. + + +2.1.1 (2026-02-03) +~~~~~~~~~~~~~~~~~~~ + +Added +^^^^^ + +* Added :mod:`isaaclab.test.mock_interfaces` module providing mock implementations for unit testing + without requiring Isaac Sim. Includes: + + * Mock assets: :class:`MockArticulation`, :class:`MockRigidObject`, :class:`MockRigidObjectCollection` + with full state tracking and property management. + * Mock sensors: :class:`MockContactSensor`, :class:`MockImu`, :class:`MockFrameTransformer` + with configurable data outputs. + * Utility classes: :class:`MockArticulationBuilder`, :class:`MockSensorBuilder`, + :class:`MockWrenchComposer` for flexible mock construction. + * Factory functions for common robot morphologies (quadruped, humanoid). + * Patching utilities and decorators for easy test injection. + 2.1.0 (2026-02-02) ~~~~~~~~~~~~~~~~~~~ diff --git a/source/isaaclab/isaaclab/devices/openxr/xr_anchor_utils.py b/source/isaaclab/isaaclab/devices/openxr/xr_anchor_utils.py index c525b99e66b..6bc9fcd6ec2 100644 --- a/source/isaaclab/isaaclab/devices/openxr/xr_anchor_utils.py +++ b/source/isaaclab/isaaclab/devices/openxr/xr_anchor_utils.py @@ -3,8 +3,6 @@ # # SPDX-License-Identifier: BSD-3-Clause -# Copyright (c) 2022-2025, The Isaac Lab Project Developers. -# SPDX-License-Identifier: BSD-3-Clause """Utilities for synchronizing XR anchor pose with a reference prim and XR config.""" from __future__ import annotations diff --git a/source/isaaclab/isaaclab/test/benchmark/__init__.py b/source/isaaclab/isaaclab/test/benchmark/__init__.py new file mode 100644 index 00000000000..4912ca1f757 --- /dev/null +++ b/source/isaaclab/isaaclab/test/benchmark/__init__.py @@ -0,0 +1,51 @@ +# 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 + +""" Benchmarking utilities for IsaacLab. + +This package provides benchmarking utilities used across different test modules. +""" + +from .benchmark_core import ( + BenchmarkConfig, + BenchmarkResult, + MethodBenchmark, + make_warp_body_mask, + make_warp_env_mask, + make_warp_joint_mask, + make_tensor_env_ids, + make_tensor_body_ids, + make_tensor_joint_ids, + benchmark_method, +) +from .benchmark_io import ( + export_results_csv, + export_results_json, + get_default_output_filename, + get_git_info, + get_hardware_info, + print_hardware_info, + print_results, +) + +__all__ = [ + "BenchmarkConfig", + "BenchmarkResult", + "MethodBenchmark", + "make_warp_body_mask", + "make_warp_env_mask", + "make_warp_joint_mask", + "make_tensor_env_ids", + "make_tensor_body_ids", + "make_tensor_joint_ids", + "export_results_csv", + "export_results_json", + "get_default_output_filename", + "get_git_info", + "get_hardware_info", + "print_hardware_info", + "print_results", + "benchmark_method", +] diff --git a/source/isaaclab/isaaclab/test/benchmark/benchmark_core.py b/source/isaaclab/isaaclab/test/benchmark/benchmark_core.py new file mode 100644 index 00000000000..ab82a1a48b1 --- /dev/null +++ b/source/isaaclab/isaaclab/test/benchmark/benchmark_core.py @@ -0,0 +1,302 @@ +# 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 + +"""Core benchmarking utilities shared across test modules. + +This module provides dataclasses, enums, and helper functions used by the +benchmark scripts for both Articulation and RigidObject classes. +""" + +from __future__ import annotations + +import contextlib +import time +from collections.abc import Callable +from dataclasses import dataclass + +import numpy as np +import warp as wp + + +@dataclass +class BenchmarkConfig: + """Configuration for the benchmarking framework.""" + + num_iterations: int = 1000 + """Number of iterations to run each function.""" + + warmup_steps: int = 10 + """Number of warmup steps before timing.""" + + num_instances: int = 4096 + """Number of instances (articulations or rigid objects).""" + + num_bodies: int = 12 + """Number of bodies per instance.""" + + num_joints: int = 11 + """Number of joints per instance (only applicable for articulations).""" + + device: str = "cuda:0" + """Device to run benchmarks on.""" + + mode: str | list[str] = "all" + """Benchmark mode(s) to run. Can be a single string or list of strings. 'all' runs all available modes.""" + + +@dataclass +class BenchmarkResult: + """Result of a single benchmark.""" + + name: str + """Name of the benchmarked method/property.""" + + mean_time_us: float + """Mean execution time in microseconds.""" + + std_time_us: float + """Standard deviation of execution time in microseconds.""" + + num_iterations: int + """Number of iterations run.""" + + mode: str | None = None + """Input mode used (e.g., 'warp', 'torch_list', etc.). None for property benchmarks.""" + + skipped: bool = False + """Whether the benchmark was skipped.""" + + skip_reason: str = "" + """Reason for skipping the benchmark.""" + + dependencies: list[str] | None = None + """List of parent properties that were pre-computed before timing.""" + + +@dataclass +class MethodBenchmark: + """Definition of a method to benchmark.""" + + name: str + """Name of the method.""" + + method_name: str + """Actual method name on the class.""" + + input_generators: dict[str, Callable] + """Dictionary mapping mode names to input generator functions.""" + + category: str = "general" + """Category of the method (e.g., 'root_state', 'joint_state', 'joint_params').""" + + +# ============================================================================= +# Common Input Generator Helpers +# ============================================================================= + +import torch + + +def make_tensor_env_ids(num_instances: int, device: str) -> torch.Tensor: + """Create a tensor of environment IDs. + + Args: + num_instances: Number of environment instances. + device: Device to create the tensor on. + + Returns: + Tensor of environment IDs [0, 1, ..., num_instances-1]. + """ + return torch.arange(num_instances, dtype=torch.long, device=device) + + +def make_tensor_joint_ids(num_joints: int, device: str) -> torch.Tensor: + """Create a tensor of joint IDs. + + Args: + num_joints: Number of joints. + device: Device to create the tensor on. + + Returns: + Tensor of joint IDs [0, 1, ..., num_joints-1]. + """ + return torch.arange(num_joints, dtype=torch.long, device=device) + + +def make_tensor_body_ids(num_bodies: int, device: str) -> torch.Tensor: + """Create a tensor of body IDs. + + Args: + num_bodies: Number of bodies. + device: Device to create the tensor on. + + Returns: + Tensor of body IDs [0, 1, ..., num_bodies-1]. + """ + return torch.arange(num_bodies, dtype=torch.long, device=device) + + +def make_warp_env_mask(num_instances: int, device: str) -> wp.array: + """Create an all-true environment mask. + + Args: + num_instances: Number of environment instances. + device: Device to create the mask on. + + Returns: + Warp array of booleans, all set to True. + """ + return wp.ones((num_instances,), dtype=wp.bool, device=device) + + +def make_warp_joint_mask(num_joints: int, device: str) -> wp.array: + """Create an all-true joint mask. + + Args: + num_joints: Number of joints. + device: Device to create the mask on. + + Returns: + Warp array of booleans, all set to True. + """ + return wp.ones((num_joints,), dtype=wp.bool, device=device) + + +def make_warp_body_mask(num_bodies: int, device: str) -> wp.array: + """Create an all-true body mask. + + Args: + num_bodies: Number of bodies. + device: Device to create the mask on. + + Returns: + Warp array of booleans, all set to True. + """ + return wp.ones((num_bodies,), dtype=wp.bool, device=device) + + +# ============================================================================= +# Benchmark Method Helper Functions +# ============================================================================= + + +def benchmark_method( + method: Callable | None, + method_name: str, + generator: Callable, + config: BenchmarkConfig, + dependencies: dict[str, list[str]] = {}, +) -> BenchmarkResult: + """Benchmarks a method. + + Args: + method: The method to benchmark. + method_name: The name of the method. + generator: The input generator to use for the method. + config: Benchmark configuration. + + Returns: + BenchmarkResult with timing statistics. + """ + # Check if method exists + if method is None: + return BenchmarkResult( + name=method_name, + mean_time_us=0.0, + std_time_us=0.0, + num_iterations=0, + skipped=True, + skip_reason="Method not found", + ) + + # Try to access the property once to check if it raises NotImplementedError + try: + _ = method(**generator(config)) + except NotImplementedError as e: + return BenchmarkResult( + name=method_name, + mean_time_us=0.0, + std_time_us=0.0, + num_iterations=0, + skipped=True, + skip_reason=f"NotImplementedError: {e}", + ) + except Exception as e: + return BenchmarkResult( + name=method_name, + mean_time_us=0.0, + std_time_us=0.0, + num_iterations=0, + skipped=True, + skip_reason=f"Error: {type(e).__name__}: {e}", + ) + + # Get dependencies for this property (if any) + dependencies_ = dependencies.get(method_name, []) + + # Warmup phase with random data + for _ in range(config.warmup_steps): + try: + # Warm up dependencies first + inputs = generator(config) + for dep in dependencies_: + _ = method(**inputs) + # Then warm up the method + _ = method(**inputs) + except Exception: + pass + # Sync GPU + if config.device.startswith("cuda"): + wp.synchronize() + + # Timing phase + times = [] + for _ in range(config.num_iterations): + # Call dependencies first to populate their caches (not timed) + # This ensures we only measure the overhead of the derived property + inputs = generator(config) + with contextlib.suppress(Exception): + for dep in dependencies_: + _ = method(**inputs) + + # Sync before timing + if config.device.startswith("cuda"): + wp.synchronize() + if torch.cuda.is_available(): + torch.cuda.synchronize() + + # Time only the target property access + start_time = time.perf_counter() + try: + _ = method(**inputs) + except Exception: + continue + + # Sync after to ensure kernel completion + if config.device.startswith("cuda"): + wp.synchronize() + if torch.cuda.is_available(): + torch.cuda.synchronize() + + end_time = time.perf_counter() + times.append((end_time - start_time) * 1e6) # Convert to microseconds + + if not times: + return BenchmarkResult( + name=method_name, + mean_time_us=0.0, + std_time_us=0.0, + num_iterations=0, + skipped=True, + skip_reason="No successful iterations", + ) + + return BenchmarkResult( + name=method_name, + mean_time_us=float(np.mean(times)), + std_time_us=float(np.std(times)), + num_iterations=len(times), + dependencies=dependencies_ if dependencies_ else None, + ) diff --git a/source/isaaclab/isaaclab/test/benchmark/benchmark_io.py b/source/isaaclab/isaaclab/test/benchmark/benchmark_io.py new file mode 100644 index 00000000000..1a4023443aa --- /dev/null +++ b/source/isaaclab/isaaclab/test/benchmark/benchmark_io.py @@ -0,0 +1,560 @@ +# 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 + +"""Benchmark I/O utilities for gathering system info and exporting results. + +This module provides functions for: +- Gathering git repository information +- Gathering hardware information (CPU, GPU, memory) +- Printing hardware information to console +- Exporting benchmark results to JSON/CSV +- Printing formatted benchmark result tables +""" + +from __future__ import annotations + +import contextlib +import json +import os +import platform +import subprocess +from datetime import datetime +from typing import TYPE_CHECKING + +import numpy as np +import torch +import warp as wp + +if TYPE_CHECKING: + from .benchmark_core import BenchmarkConfig, BenchmarkResult + + +def get_git_info() -> dict: + """Get git repository information. + + Returns: + Dictionary containing git commit hash, branch, and other info. + """ + git_info = { + "commit_hash": "Unknown", + "commit_hash_short": "Unknown", + "branch": "Unknown", + "commit_date": "Unknown", + } + + # Get the directory of the caller to find the repo root + script_dir = os.path.dirname(os.path.abspath(__file__)) + + try: + # Get full commit hash + result = subprocess.run( + ["git", "rev-parse", "HEAD"], + cwd=script_dir, + capture_output=True, + text=True, + timeout=5, + ) + if result.returncode == 0: + git_info["commit_hash"] = result.stdout.strip() + git_info["commit_hash_short"] = result.stdout.strip()[:8] + + # Get branch name + result = subprocess.run( + ["git", "rev-parse", "--abbrev-ref", "HEAD"], + cwd=script_dir, + capture_output=True, + text=True, + timeout=5, + ) + if result.returncode == 0: + git_info["branch"] = result.stdout.strip() + + # Get commit date + result = subprocess.run( + ["git", "log", "-1", "--format=%ci"], + cwd=script_dir, + capture_output=True, + text=True, + timeout=5, + ) + if result.returncode == 0: + git_info["commit_date"] = result.stdout.strip() + + except Exception: + pass + + return git_info + + +def get_hardware_info() -> dict: + """Gather hardware information for the benchmark. + + Returns: + Dictionary containing CPU, GPU, and memory information. + """ + hardware_info = { + "cpu": { + "name": platform.processor() or "Unknown", + "physical_cores": os.cpu_count(), + }, + "gpu": {}, + "memory": {}, + "system": { + "platform": platform.system(), + "platform_release": platform.release(), + "platform_version": platform.version(), + "architecture": platform.machine(), + "python_version": platform.python_version(), + }, + } + + # Try to get more detailed CPU info on Linux + with contextlib.suppress(Exception): + with open("/proc/cpuinfo") as f: + cpuinfo = f.read() + for line in cpuinfo.split("\n"): + if "model name" in line: + hardware_info["cpu"]["name"] = line.split(":")[1].strip() + break + + # Memory info + try: + with open("/proc/meminfo") as f: + meminfo = f.read() + for line in meminfo.split("\n"): + if "MemTotal" in line: + # Convert from KB to GB + mem_kb = int(line.split()[1]) + hardware_info["memory"]["total_gb"] = round(mem_kb / (1024 * 1024), 2) + break + except Exception: + # Fallback using psutil if available + try: + import psutil + + mem = psutil.virtual_memory() + hardware_info["memory"]["total_gb"] = round(mem.total / (1024**3), 2) + except ImportError: + hardware_info["memory"]["total_gb"] = "Unknown" + + # GPU info using PyTorch + if torch.cuda.is_available(): + hardware_info["gpu"]["available"] = True + hardware_info["gpu"]["count"] = torch.cuda.device_count() + hardware_info["gpu"]["devices"] = [] + + for i in range(torch.cuda.device_count()): + gpu_props = torch.cuda.get_device_properties(i) + hardware_info["gpu"]["devices"].append( + { + "index": i, + "name": gpu_props.name, + "total_memory_gb": round(gpu_props.total_memory / (1024**3), 2), + "compute_capability": f"{gpu_props.major}.{gpu_props.minor}", + "multi_processor_count": gpu_props.multi_processor_count, + } + ) + + # Current device info + current_device = torch.cuda.current_device() + hardware_info["gpu"]["current_device"] = current_device + hardware_info["gpu"]["current_device_name"] = torch.cuda.get_device_name(current_device) + else: + hardware_info["gpu"]["available"] = False + + # PyTorch and CUDA versions + hardware_info["gpu"]["pytorch_version"] = torch.__version__ + if torch.cuda.is_available(): + try: + import torch.version as torch_version + + cuda_version = getattr(torch_version, "cuda", None) + hardware_info["gpu"]["cuda_version"] = cuda_version if cuda_version else "Unknown" + except Exception: + hardware_info["gpu"]["cuda_version"] = "Unknown" + else: + hardware_info["gpu"]["cuda_version"] = "N/A" + + # Warp info + try: + warp_version = getattr(wp.config, "version", None) + hardware_info["warp"] = {"version": warp_version if warp_version else "Unknown"} + except Exception: + hardware_info["warp"] = {"version": "Unknown"} + + return hardware_info + + +def print_hardware_info(hardware_info: dict): + """Print hardware information to console. + + Args: + hardware_info: Dictionary containing hardware information. + """ + print("\n" + "=" * 80) + print("HARDWARE INFORMATION") + print("=" * 80) + + # System + sys_info = hardware_info.get("system", {}) + print(f"\nSystem: {sys_info.get('platform', 'Unknown')} {sys_info.get('platform_release', '')}") + print(f"Python: {sys_info.get('python_version', 'Unknown')}") + + # CPU + cpu_info = hardware_info.get("cpu", {}) + print(f"\nCPU: {cpu_info.get('name', 'Unknown')}") + print(f" Cores: {cpu_info.get('physical_cores', 'Unknown')}") + + # Memory + mem_info = hardware_info.get("memory", {}) + print(f"\nMemory: {mem_info.get('total_gb', 'Unknown')} GB") + + # GPU + gpu_info = hardware_info.get("gpu", {}) + if gpu_info.get("available", False): + print(f"\nGPU: {gpu_info.get('current_device_name', 'Unknown')}") + for device in gpu_info.get("devices", []): + print(f" [{device['index']}] {device['name']}") + print(f" Memory: {device['total_memory_gb']} GB") + print(f" Compute Capability: {device['compute_capability']}") + print(f" SM Count: {device['multi_processor_count']}") + print(f"\nPyTorch: {gpu_info.get('pytorch_version', 'Unknown')}") + print(f"CUDA: {gpu_info.get('cuda_version', 'Unknown')}") + else: + print("\nGPU: Not available") + + warp_info = hardware_info.get("warp", {}) + print(f"Warp: {warp_info.get('version', 'Unknown')}") + + # Repository info + repo_info = get_git_info() + print("\nRepository:") + print(f" Commit: {repo_info.get('commit_hash_short', 'Unknown')}") + print(f" Branch: {repo_info.get('branch', 'Unknown')}") + print(f" Date: {repo_info.get('commit_date', 'Unknown')}") + print("=" * 80) + + +def print_results(results: list[BenchmarkResult], include_mode: bool = True): + """Print benchmark results in a formatted table. + + Args: + results: List of BenchmarkResults to print. + include_mode: Whether to separate and compare results by input mode. + """ + print("\n" + "=" * 100) + print("BENCHMARK RESULTS") + print("=" * 100) + + # Separate skipped and completed results + completed = [r for r in results if not r.skipped] + skipped = [r for r in results if r.skipped] + + # Get unique modes present in the results + modes = sorted(list({r.mode for r in completed if r.mode is not None})) + + if include_mode and modes: + # Separate by mode + results_by_mode = {mode: [r for r in completed if r.mode == mode] for mode in modes} + + # Print comparison table if multiple modes + if len(modes) > 1: + print("\n" + "-" * 120) + mode_headers = " vs ".join([m.capitalize() for m in modes]) + print(f"COMPARISON: {mode_headers}") + print("-" * 120) + + # Use the first mode as baseline (usually 'warp') + baseline_mode = modes[0] + + header = f"{'Method Name':<35}" + for mode in modes: + header += f" {mode.capitalize()[:10]:<12}" + + # Add slowdown columns for non-baseline modes + for mode in modes[1:]: + header += f" {mode.capitalize()[:6]} Slow.<14" + + print(header) + print("-" * 120) + + results_by_name = {} + for mode in modes: + for r in results_by_mode[mode]: + if r.name not in results_by_name: + results_by_name[r.name] = {} + results_by_name[r.name][mode] = r + + for name in sorted(results_by_name.keys()): + row_results = results_by_name[name] + if baseline_mode not in row_results: + continue # Skip if baseline missing + + baseline_time = row_results[baseline_mode].mean_time_us + + row_str = f"{name:<35} {baseline_time:>9.2f}" + + # Time columns + for mode in modes[1:]: + if mode in row_results: + row_str += f" {row_results[mode].mean_time_us:>9.2f}" + else: + row_str += " " + "N/A".rjust(9) + + # Slowdown columns + for mode in modes[1:]: + if mode in row_results and baseline_time > 0: + slowdown = row_results[mode].mean_time_us / baseline_time + row_str += f" {slowdown:>10.2f}x" + else: + row_str += " " + "N/A".rjust(11) + + print(row_str) + + # Print individual results by mode + for mode in modes: + mode_results = results_by_mode[mode] + if mode_results: + print("\n" + "-" * 100) + print(f"{mode.upper()}") + print("-" * 100) + + # Sort by mean time (descending) + mode_results_sorted = sorted(mode_results, key=lambda x: x.mean_time_us, reverse=True) + + print(f"{'Method Name':<45} {'Mean (µs)':<15} {'Std (µs)':<15} {'Iterations':<12}") + print("-" * 87) + + for result in mode_results_sorted: + print( + f"{result.name:<45} {result.mean_time_us:>12.2f} {result.std_time_us:>12.2f} " + f" {result.num_iterations:>10}" + ) + + # Summary stats + mean_times = [r.mean_time_us for r in mode_results_sorted] + print("-" * 87) + if mean_times: + print(f" Fastest: {min(mean_times):.2f} µs ({mode_results_sorted[-1].name})") + print(f" Slowest: {max(mean_times):.2f} µs ({mode_results_sorted[0].name})") + print(f" Average: {np.mean(mean_times):.2f} µs") + else: + # No mode separation - just print all results + completed_sorted = sorted(completed, key=lambda x: x.mean_time_us, reverse=True) + + print(f"\n{'Property Name':<45} {'Mean (µs)':<15} {'Std (µs)':<15} {'Iterations':<12}") + print("-" * 87) + + for result in completed_sorted: + name_display = result.name + if result.dependencies: + name_display = f"{result.name} *" + print( + f"{name_display:<45} {result.mean_time_us:>12.2f} {result.std_time_us:>12.2f} " + f" {result.num_iterations:>10}" + ) + + # Print summary statistics + if completed_sorted: + print("-" * 87) + mean_times = [r.mean_time_us for r in completed_sorted] + print("\nSummary Statistics:") + print(f" Total benchmarked: {len(completed_sorted)}") + print(f" Fastest: {min(mean_times):.2f} µs ({completed_sorted[-1].name})") + print(f" Slowest: {max(mean_times):.2f} µs ({completed_sorted[0].name})") + print(f" Average: {np.mean(mean_times):.2f} µs") + print(f" Median: {np.median(mean_times):.2f} µs") + + # Print note about derived properties + derived_count = sum(1 for r in completed_sorted if r.dependencies) + if derived_count > 0: + print(f"\n * = Derived property ({derived_count} total). Dependencies were pre-computed") + print(" before timing to measure isolated overhead.") + + # Print skipped results + if skipped: + print(f"\nSkipped ({len(skipped)}):") + for result in skipped: + mode_str = f" [{result.mode}]" if result.mode else "" + print(f" - {result.name}{mode_str}: {result.skip_reason}") + + +def export_results_json( + results: list[BenchmarkResult], + config: BenchmarkConfig, + hardware_info: dict, + filename: str, + include_mode: bool = True, +): + """Export benchmark results to a JSON file. + + Args: + results: List of BenchmarkResults to export. + config: Benchmark configuration used. + hardware_info: Hardware information dictionary. + filename: Output JSON filename. + include_mode: Whether to include mode separation in output. + """ + completed = [r for r in results if not r.skipped] + skipped = [r for r in results if r.skipped] + + git_info = get_git_info() + + # Build config dict - only include relevant fields + config_dict = { + "num_iterations": config.num_iterations, + "warmup_steps": config.warmup_steps, + "num_instances": config.num_instances, + "num_bodies": config.num_bodies, + "device": config.device, + } + if hasattr(config, "num_joints") and config.num_joints > 0: + config_dict["num_joints"] = config.num_joints + if hasattr(config, "mode"): + config_dict["mode"] = config.mode + + output = { + "metadata": { + "timestamp": datetime.now().isoformat(), + "repository": git_info, + "config": config_dict, + "hardware": hardware_info, + "total_benchmarks": len(results), + "completed_benchmarks": len(completed), + "skipped_benchmarks": len(skipped), + }, + "results": {}, + "comparison": [], + "skipped": [], + } + + # Get unique modes present in the results + modes = sorted(list({r.mode for r in completed if r.mode is not None})) + + if include_mode and modes: + for mode in modes: + output["results"][mode] = [] + + for result in completed: + result_entry = { + "name": result.name, + "mean_time_us": result.mean_time_us, + "std_time_us": result.std_time_us, + "num_iterations": result.num_iterations, + } + if result.dependencies: + result_entry["dependencies"] = result.dependencies + + if result.mode in output["results"]: + output["results"][result.mode].append(result_entry) + + # Add comparison data + results_by_name = {} + for r in completed: + if r.mode is not None: + if r.name not in results_by_name: + results_by_name[r.name] = {} + results_by_name[r.name][r.mode] = r + + baseline_mode = modes[0] # Usually 'warp' or fastest + + for name in results_by_name: + if baseline_mode in results_by_name[name]: + baseline_time = results_by_name[name][baseline_mode].mean_time_us + comparison_entry = { + "name": name, + f"{baseline_mode}_time_us": baseline_time, + } + + for mode in modes: + if mode == baseline_mode: + continue + if mode in results_by_name[name]: + mode_time = results_by_name[name][mode].mean_time_us + comparison_entry[f"{mode}_time_us"] = mode_time + comparison_entry[f"{mode}_slowdown"] = mode_time / baseline_time if baseline_time > 0 else None + + output["comparison"].append(comparison_entry) + else: + output["results"] = [] + for result in completed: + result_entry = { + "name": result.name, + "mean_time_us": result.mean_time_us, + "std_time_us": result.std_time_us, + "num_iterations": result.num_iterations, + } + if result.dependencies: + result_entry["dependencies"] = result.dependencies + result_entry["note"] = "Timing excludes dependency computation (dependencies were pre-computed)" + output["results"].append(result_entry) + + # Add skipped + for result in skipped: + entry = {"name": result.name, "reason": result.skip_reason} + if result.mode: + entry["mode"] = result.mode + output["skipped"].append(entry) + + with open(filename, "w") as jsonfile: + json.dump(output, jsonfile, indent=2) + + print(f"\nResults exported to {filename}") + + +def export_results_csv(results: list[BenchmarkResult], filename: str): + """Export benchmark results to a CSV file. + + Args: + results: List of BenchmarkResults to export. + filename: Output CSV filename. + """ + import csv + + with open(filename, "w", newline="") as csvfile: + writer = csv.writer(csvfile) + writer.writerow( + [ + "Name", + "Mode", + "Mean (µs)", + "Std (µs)", + "Iterations", + "Dependencies", + "Skipped", + "Skip Reason", + ] + ) + + for result in results: + deps_str = ", ".join(result.dependencies) if result.dependencies else "" + mode_str = result.mode if result.mode else "" + writer.writerow( + [ + result.name, + mode_str, + f"{result.mean_time_us:.4f}" if not result.skipped else "", + f"{result.std_time_us:.4f}" if not result.skipped else "", + result.num_iterations, + deps_str, + result.skipped, + result.skip_reason, + ] + ) + + print(f"\nResults exported to {filename}") + + +def get_default_output_filename(prefix: str = "benchmark") -> str: + """Generate default output filename with current date and time. + + Args: + prefix: Prefix for the filename (e.g., "articulation_benchmark"). + + Returns: + Filename string with timestamp. + """ + datetime_str = datetime.now().strftime("%Y-%m-%d_%H-%M-%S") + return f"{prefix}_{datetime_str}.json" diff --git a/source/isaaclab/isaaclab/test/mock_interfaces/__init__.py b/source/isaaclab/isaaclab/test/mock_interfaces/__init__.py new file mode 100644 index 00000000000..8c3fe7595a1 --- /dev/null +++ b/source/isaaclab/isaaclab/test/mock_interfaces/__init__.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 + +"""Mock interfaces for IsaacLab sensors and assets. + +This module provides mock implementations of sensor and asset classes for unit testing +without requiring the full Isaac Sim simulation environment. + +Example usage: + + .. code-block:: python + + from isaaclab.test.mock_interfaces.sensors import MockContactSensor + from isaaclab.test.mock_interfaces.assets import MockArticulation + + # Create mock sensor + sensor = MockContactSensor(num_instances=4, num_bodies=4, device="cpu") + sensor.data.set_mock_data(net_forces_w=torch.randn(4, 4, 3)) + + # Create mock articulation + robot = MockArticulation( + num_instances=4, + num_joints=12, + num_bodies=13, + joint_names=["joint_" + str(i) for i in range(12)], + device="cpu" + ) + +""" + +from .assets import * +from .sensors import * +from .utils import * diff --git a/source/isaaclab/isaaclab/test/mock_interfaces/assets/__init__.py b/source/isaaclab/isaaclab/test/mock_interfaces/assets/__init__.py new file mode 100644 index 00000000000..f6169a1fcc7 --- /dev/null +++ b/source/isaaclab/isaaclab/test/mock_interfaces/assets/__init__.py @@ -0,0 +1,17 @@ +# 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 + +"""Mock asset interfaces for testing without Isaac Sim.""" + +from .mock_articulation import MockArticulation, MockArticulationData +from .mock_rigid_object import MockRigidObject, MockRigidObjectData +from .mock_rigid_object_collection import MockRigidObjectCollection, MockRigidObjectCollectionData +from .factories import ( + create_mock_articulation, + create_mock_humanoid, + create_mock_quadruped, + create_mock_rigid_object, + create_mock_rigid_object_collection, +) diff --git a/source/isaaclab/isaaclab/test/mock_interfaces/assets/factories.py b/source/isaaclab/isaaclab/test/mock_interfaces/assets/factories.py new file mode 100644 index 00000000000..c135c80005f --- /dev/null +++ b/source/isaaclab/isaaclab/test/mock_interfaces/assets/factories.py @@ -0,0 +1,209 @@ +# 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 + +"""Factory functions for creating pre-configured mock assets.""" + +from __future__ import annotations + +import torch + +from .mock_articulation import MockArticulation +from .mock_rigid_object import MockRigidObject +from .mock_rigid_object_collection import MockRigidObjectCollection + + +def create_mock_articulation( + num_instances: int = 1, + num_joints: int = 1, + num_bodies: int = 2, + joint_names: list[str] | None = None, + body_names: list[str] | None = None, + is_fixed_base: bool = False, + device: str = "cpu", +) -> MockArticulation: + """Create a mock articulation with default configuration. + + Args: + num_instances: Number of articulation instances. + num_joints: Number of joints. + num_bodies: Number of bodies. + joint_names: Names of joints. + body_names: Names of bodies. + is_fixed_base: Whether the articulation has a fixed base. + device: Device for tensor allocation. + + Returns: + Configured MockArticulation instance. + """ + return MockArticulation( + num_instances=num_instances, + num_joints=num_joints, + num_bodies=num_bodies, + joint_names=joint_names, + body_names=body_names, + is_fixed_base=is_fixed_base, + device=device, + ) + + +def create_mock_quadruped( + num_instances: int = 1, + device: str = "cpu", +) -> MockArticulation: + """Create a mock quadruped robot articulation. + + Creates a quadruped with 12 joints (3 per leg) and 13 bodies. + + Args: + num_instances: Number of robot instances. + device: Device for tensor allocation. + + Returns: + Configured MockArticulation instance for a quadruped. + """ + leg_names = ["FL", "FR", "RL", "RR"] + joint_suffixes = ["hip", "thigh", "calf"] + + joint_names = [f"{leg}_{suffix}" for leg in leg_names for suffix in joint_suffixes] + body_names = ["base"] + [f"{leg}_{part}" for leg in leg_names for part in ["hip", "thigh", "calf"]] + + robot = MockArticulation( + num_instances=num_instances, + num_joints=12, + num_bodies=13, + joint_names=joint_names, + body_names=body_names, + is_fixed_base=False, + device=device, + ) + + # Set reasonable default joint limits for a quadruped + joint_pos_limits = torch.zeros(num_instances, 12, 2, device=device) + joint_pos_limits[..., 0] = -1.57 # Lower limit + joint_pos_limits[..., 1] = 1.57 # Upper limit + robot.data.set_joint_pos_limits(joint_pos_limits) + + return robot + + +def create_mock_humanoid( + num_instances: int = 1, + device: str = "cpu", +) -> MockArticulation: + """Create a mock humanoid robot articulation. + + Creates a humanoid with 21 joints and 22 bodies. + + Args: + num_instances: Number of robot instances. + device: Device for tensor allocation. + + Returns: + Configured MockArticulation instance for a humanoid. + """ + # Simplified humanoid joint structure + joint_names = [ + # Torso + "torso_yaw", + "torso_pitch", + "torso_roll", + # Left arm + "L_shoulder_pitch", + "L_shoulder_roll", + "L_shoulder_yaw", + "L_elbow", + # Right arm + "R_shoulder_pitch", + "R_shoulder_roll", + "R_shoulder_yaw", + "R_elbow", + # Left leg + "L_hip_yaw", + "L_hip_roll", + "L_hip_pitch", + "L_knee", + "L_ankle_pitch", + # Right leg + "R_hip_yaw", + "R_hip_roll", + "R_hip_pitch", + "R_knee", + "R_ankle_pitch", + ] + + body_names = [ + "pelvis", + "torso", + "L_upper_arm", + "L_lower_arm", + "L_hand", + "R_upper_arm", + "R_lower_arm", + "R_hand", + "L_thigh", + "L_shin", + "L_foot", + "R_thigh", + "R_shin", + "R_foot", + "head", + ] + + return MockArticulation( + num_instances=num_instances, + num_joints=21, + num_bodies=len(body_names), + joint_names=joint_names, + body_names=body_names, + is_fixed_base=False, + device=device, + ) + + +def create_mock_rigid_object( + num_instances: int = 1, + body_names: list[str] | None = None, + device: str = "cpu", +) -> MockRigidObject: + """Create a mock rigid object with default configuration. + + Args: + num_instances: Number of rigid object instances. + body_names: Names of bodies. + device: Device for tensor allocation. + + Returns: + Configured MockRigidObject instance. + """ + return MockRigidObject( + num_instances=num_instances, + body_names=body_names, + device=device, + ) + + +def create_mock_rigid_object_collection( + num_instances: int = 1, + num_bodies: int = 1, + body_names: list[str] | None = None, + device: str = "cpu", +) -> MockRigidObjectCollection: + """Create a mock rigid object collection with default configuration. + + Args: + num_instances: Number of environment instances. + num_bodies: Number of bodies in the collection. + body_names: Names of bodies. + device: Device for tensor allocation. + + Returns: + Configured MockRigidObjectCollection instance. + """ + return MockRigidObjectCollection( + num_instances=num_instances, + num_bodies=num_bodies, + body_names=body_names, + device=device, + ) diff --git a/source/isaaclab/isaaclab/test/mock_interfaces/assets/mock_articulation.py b/source/isaaclab/isaaclab/test/mock_interfaces/assets/mock_articulation.py new file mode 100644 index 00000000000..7004da6798b --- /dev/null +++ b/source/isaaclab/isaaclab/test/mock_interfaces/assets/mock_articulation.py @@ -0,0 +1,1609 @@ +# 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 + +"""Mock articulation asset for testing without Isaac Sim.""" + +from __future__ import annotations + +import re +from collections.abc import Sequence + +import torch + + +class MockArticulationData: + """Mock data container for articulation asset. + + This class mimics the interface of BaseArticulationData for testing purposes. + All tensor properties return zero tensors with correct shapes if not explicitly set. + """ + + def __init__( + self, + num_instances: int, + num_joints: int, + num_bodies: int, + joint_names: list[str] | None = None, + body_names: list[str] | None = None, + num_fixed_tendons: int = 0, + num_spatial_tendons: int = 0, + fixed_tendon_names: list[str] | None = None, + spatial_tendon_names: list[str] | None = None, + device: str = "cpu", + ): + """Initialize mock articulation data. + + Args: + num_instances: Number of articulation instances. + num_joints: Number of joints. + num_bodies: Number of bodies. + joint_names: Names of joints. + body_names: Names of bodies. + num_fixed_tendons: Number of fixed tendons. + num_spatial_tendons: Number of spatial tendons. + fixed_tendon_names: Names of fixed tendons. + spatial_tendon_names: Names of spatial tendons. + device: Device for tensor allocation. + """ + self._num_instances = num_instances + self._num_joints = num_joints + self._num_bodies = num_bodies + self._num_fixed_tendons = num_fixed_tendons + self._num_spatial_tendons = num_spatial_tendons + self._device = device + + # Names + self.joint_names = joint_names or [f"joint_{i}" for i in range(num_joints)] + self.body_names = body_names or [f"body_{i}" for i in range(num_bodies)] + self.fixed_tendon_names = fixed_tendon_names or [f"fixed_tendon_{i}" for i in range(num_fixed_tendons)] + self.spatial_tendon_names = spatial_tendon_names or [f"spatial_tendon_{i}" for i in range(num_spatial_tendons)] + + # -- Internal storage for mock data -- + # Default states + self._default_root_pose: torch.Tensor | None = None + self._default_root_vel: torch.Tensor | None = None + self._default_joint_pos: torch.Tensor | None = None + self._default_joint_vel: torch.Tensor | None = None + + # Joint commands + self._joint_pos_target: torch.Tensor | None = None + self._joint_vel_target: torch.Tensor | None = None + self._joint_effort_target: torch.Tensor | None = None + self._computed_torque: torch.Tensor | None = None + self._applied_torque: torch.Tensor | None = None + + # Joint properties + self._joint_stiffness: torch.Tensor | None = None + self._joint_damping: torch.Tensor | None = None + self._joint_armature: torch.Tensor | None = None + self._joint_friction_coeff: torch.Tensor | None = None + self._joint_dynamic_friction_coeff: torch.Tensor | None = None + self._joint_viscous_friction_coeff: torch.Tensor | None = None + self._joint_pos_limits: torch.Tensor | None = None + self._joint_vel_limits: torch.Tensor | None = None + self._joint_effort_limits: torch.Tensor | None = None + self._soft_joint_pos_limits: torch.Tensor | None = None + self._soft_joint_vel_limits: torch.Tensor | None = None + self._gear_ratio: torch.Tensor | None = None + + # Joint state + self._joint_pos: torch.Tensor | None = None + self._joint_vel: torch.Tensor | None = None + self._joint_acc: torch.Tensor | None = None + + # Root state (link frame) + self._root_link_pose_w: torch.Tensor | None = None + self._root_link_vel_w: torch.Tensor | None = None + + # Root state (CoM frame) + self._root_com_pose_w: torch.Tensor | None = None + self._root_com_vel_w: torch.Tensor | None = None + + # Body state (link frame) + self._body_link_pose_w: torch.Tensor | None = None + self._body_link_vel_w: torch.Tensor | None = None + + # Body state (CoM frame) + self._body_com_pose_w: torch.Tensor | None = None + self._body_com_vel_w: torch.Tensor | None = None + self._body_com_acc_w: torch.Tensor | None = None + self._body_com_pose_b: torch.Tensor | None = None + + # Body properties + self._body_mass: torch.Tensor | None = None + self._body_inertia: torch.Tensor | None = None + self._body_incoming_joint_wrench_b: torch.Tensor | None = None + + # Tendon properties (fixed) + self._fixed_tendon_stiffness: torch.Tensor | None = None + self._fixed_tendon_damping: torch.Tensor | None = None + self._fixed_tendon_limit_stiffness: torch.Tensor | None = None + self._fixed_tendon_rest_length: torch.Tensor | None = None + self._fixed_tendon_offset: torch.Tensor | None = None + self._fixed_tendon_pos_limits: torch.Tensor | None = None + + # Tendon properties (spatial) + self._spatial_tendon_stiffness: torch.Tensor | None = None + self._spatial_tendon_damping: torch.Tensor | None = None + self._spatial_tendon_limit_stiffness: torch.Tensor | None = None + self._spatial_tendon_offset: torch.Tensor | None = None + + @property + def device(self) -> str: + """Device for tensor allocation.""" + return self._device + + # -- Helper for identity quaternion -- + def _identity_quat(self, *shape: int) -> torch.Tensor: + """Create identity quaternion tensor (w, x, y, z) = (1, 0, 0, 0).""" + quat = torch.zeros(*shape, 4, device=self._device) + quat[..., 0] = 1.0 + return quat + + # -- Default state properties -- + + @property + def default_root_pose(self) -> torch.Tensor: + """Default root pose [pos(3), quat_xyzw(4)]. Shape: (N, 7).""" + if self._default_root_pose is None: + pose = torch.zeros(self._num_instances, 7, device=self._device) + pose[:, 3:7] = self._identity_quat(self._num_instances) + return pose + return self._default_root_pose + + @property + def default_root_vel(self) -> torch.Tensor: + """Default root velocity [lin_vel(3), ang_vel(3)]. Shape: (N, 6).""" + if self._default_root_vel is None: + return torch.zeros(self._num_instances, 6, device=self._device) + return self._default_root_vel + + @property + def default_root_state(self) -> torch.Tensor: + """Default root state [pose(7), vel(6)]. Shape: (N, 13).""" + return torch.cat([self.default_root_pose, self.default_root_vel], dim=-1) + + @property + def default_joint_pos(self) -> torch.Tensor: + """Default joint positions. Shape: (N, num_joints).""" + if self._default_joint_pos is None: + return torch.zeros(self._num_instances, self._num_joints, device=self._device) + return self._default_joint_pos + + @property + def default_joint_vel(self) -> torch.Tensor: + """Default joint velocities. Shape: (N, num_joints).""" + if self._default_joint_vel is None: + return torch.zeros(self._num_instances, self._num_joints, device=self._device) + return self._default_joint_vel + + # -- Joint command properties -- + + @property + def joint_pos_target(self) -> torch.Tensor: + """Joint position targets. Shape: (N, num_joints).""" + if self._joint_pos_target is None: + return torch.zeros(self._num_instances, self._num_joints, device=self._device) + return self._joint_pos_target + + @property + def joint_vel_target(self) -> torch.Tensor: + """Joint velocity targets. Shape: (N, num_joints).""" + if self._joint_vel_target is None: + return torch.zeros(self._num_instances, self._num_joints, device=self._device) + return self._joint_vel_target + + @property + def joint_effort_target(self) -> torch.Tensor: + """Joint effort targets. Shape: (N, num_joints).""" + if self._joint_effort_target is None: + return torch.zeros(self._num_instances, self._num_joints, device=self._device) + return self._joint_effort_target + + @property + def computed_torque(self) -> torch.Tensor: + """Computed torques before clipping. Shape: (N, num_joints).""" + if self._computed_torque is None: + return torch.zeros(self._num_instances, self._num_joints, device=self._device) + return self._computed_torque + + @property + def applied_torque(self) -> torch.Tensor: + """Applied torques after clipping. Shape: (N, num_joints).""" + if self._applied_torque is None: + return torch.zeros(self._num_instances, self._num_joints, device=self._device) + return self._applied_torque + + # -- Joint properties -- + + @property + def joint_stiffness(self) -> torch.Tensor: + """Joint stiffness. Shape: (N, num_joints).""" + if self._joint_stiffness is None: + return torch.zeros(self._num_instances, self._num_joints, device=self._device) + return self._joint_stiffness + + @property + def joint_damping(self) -> torch.Tensor: + """Joint damping. Shape: (N, num_joints).""" + if self._joint_damping is None: + return torch.zeros(self._num_instances, self._num_joints, device=self._device) + return self._joint_damping + + @property + def joint_armature(self) -> torch.Tensor: + """Joint armature. Shape: (N, num_joints).""" + if self._joint_armature is None: + return torch.zeros(self._num_instances, self._num_joints, device=self._device) + return self._joint_armature + + @property + def joint_friction_coeff(self) -> torch.Tensor: + """Joint static friction coefficient. Shape: (N, num_joints).""" + if self._joint_friction_coeff is None: + return torch.zeros(self._num_instances, self._num_joints, device=self._device) + return self._joint_friction_coeff + + @property + def joint_dynamic_friction_coeff(self) -> torch.Tensor: + """Joint dynamic friction coefficient. Shape: (N, num_joints).""" + if self._joint_dynamic_friction_coeff is None: + return torch.zeros(self._num_instances, self._num_joints, device=self._device) + return self._joint_dynamic_friction_coeff + + @property + def joint_viscous_friction_coeff(self) -> torch.Tensor: + """Joint viscous friction coefficient. Shape: (N, num_joints).""" + if self._joint_viscous_friction_coeff is None: + return torch.zeros(self._num_instances, self._num_joints, device=self._device) + return self._joint_viscous_friction_coeff + + @property + def joint_pos_limits(self) -> torch.Tensor: + """Joint position limits [lower, upper]. Shape: (N, num_joints, 2).""" + if self._joint_pos_limits is None: + limits = torch.zeros(self._num_instances, self._num_joints, 2, device=self._device) + limits[..., 0] = -float("inf") + limits[..., 1] = float("inf") + return limits + return self._joint_pos_limits + + @property + def joint_vel_limits(self) -> torch.Tensor: + """Joint velocity limits. Shape: (N, num_joints).""" + if self._joint_vel_limits is None: + return torch.full((self._num_instances, self._num_joints), float("inf"), device=self._device) + return self._joint_vel_limits + + @property + def joint_effort_limits(self) -> torch.Tensor: + """Joint effort limits. Shape: (N, num_joints).""" + if self._joint_effort_limits is None: + return torch.full((self._num_instances, self._num_joints), float("inf"), device=self._device) + return self._joint_effort_limits + + @property + def soft_joint_pos_limits(self) -> torch.Tensor: + """Soft joint position limits. Shape: (N, num_joints, 2).""" + if self._soft_joint_pos_limits is None: + return self.joint_pos_limits.clone() + return self._soft_joint_pos_limits + + @property + def soft_joint_vel_limits(self) -> torch.Tensor: + """Soft joint velocity limits. Shape: (N, num_joints).""" + if self._soft_joint_vel_limits is None: + return self.joint_vel_limits.clone() + return self._soft_joint_vel_limits + + @property + def gear_ratio(self) -> torch.Tensor: + """Gear ratio. Shape: (N, num_joints).""" + if self._gear_ratio is None: + return torch.ones(self._num_instances, self._num_joints, device=self._device) + return self._gear_ratio + + # -- Joint state properties -- + + @property + def joint_pos(self) -> torch.Tensor: + """Joint positions. Shape: (N, num_joints).""" + if self._joint_pos is None: + return torch.zeros(self._num_instances, self._num_joints, device=self._device) + return self._joint_pos + + @property + def joint_vel(self) -> torch.Tensor: + """Joint velocities. Shape: (N, num_joints).""" + if self._joint_vel is None: + return torch.zeros(self._num_instances, self._num_joints, device=self._device) + return self._joint_vel + + @property + def joint_acc(self) -> torch.Tensor: + """Joint accelerations. Shape: (N, num_joints).""" + if self._joint_acc is None: + return torch.zeros(self._num_instances, self._num_joints, device=self._device) + return self._joint_acc + + # -- Root state properties (link frame) -- + + @property + def root_link_pose_w(self) -> torch.Tensor: + """Root link pose in world frame. Shape: (N, 7).""" + if self._root_link_pose_w is None: + pose = torch.zeros(self._num_instances, 7, device=self._device) + pose[:, 3:7] = self._identity_quat(self._num_instances) + return pose + return self._root_link_pose_w + + @property + def root_link_vel_w(self) -> torch.Tensor: + """Root link velocity in world frame. Shape: (N, 6).""" + if self._root_link_vel_w is None: + return torch.zeros(self._num_instances, 6, device=self._device) + return self._root_link_vel_w + + @property + def root_link_state_w(self) -> torch.Tensor: + """Root link state in world frame. Shape: (N, 13).""" + return torch.cat([self.root_link_pose_w, self.root_link_vel_w], dim=-1) + + # Sliced properties + @property + def root_link_pos_w(self) -> torch.Tensor: + """Root link position in world frame. Shape: (N, 3).""" + return self.root_link_pose_w[:, :3] + + @property + def root_link_quat_w(self) -> torch.Tensor: + """Root link orientation in world frame. Shape: (N, 4).""" + return self.root_link_pose_w[:, 3:7] + + @property + def root_link_lin_vel_w(self) -> torch.Tensor: + """Root link linear velocity in world frame. Shape: (N, 3).""" + return self.root_link_vel_w[:, :3] + + @property + def root_link_ang_vel_w(self) -> torch.Tensor: + """Root link angular velocity in world frame. Shape: (N, 3).""" + return self.root_link_vel_w[:, 3:6] + + # -- Root state properties (CoM frame) -- + + @property + def root_com_pose_w(self) -> torch.Tensor: + """Root CoM pose in world frame. Shape: (N, 7).""" + if self._root_com_pose_w is None: + return self.root_link_pose_w.clone() + return self._root_com_pose_w + + @property + def root_com_vel_w(self) -> torch.Tensor: + """Root CoM velocity in world frame. Shape: (N, 6).""" + if self._root_com_vel_w is None: + return self.root_link_vel_w.clone() + return self._root_com_vel_w + + @property + def root_com_state_w(self) -> torch.Tensor: + """Root CoM state in world frame. Shape: (N, 13).""" + return torch.cat([self.root_com_pose_w, self.root_com_vel_w], dim=-1) + + @property + def root_state_w(self) -> torch.Tensor: + """Root state (link pose + CoM velocity). Shape: (N, 13).""" + return torch.cat([self.root_link_pose_w, self.root_com_vel_w], dim=-1) + + # Sliced properties + @property + def root_com_pos_w(self) -> torch.Tensor: + """Root CoM position in world frame. Shape: (N, 3).""" + return self.root_com_pose_w[:, :3] + + @property + def root_com_quat_w(self) -> torch.Tensor: + """Root CoM orientation in world frame. Shape: (N, 4).""" + return self.root_com_pose_w[:, 3:7] + + @property + def root_com_lin_vel_w(self) -> torch.Tensor: + """Root CoM linear velocity in world frame. Shape: (N, 3).""" + return self.root_com_vel_w[:, :3] + + @property + def root_com_ang_vel_w(self) -> torch.Tensor: + """Root CoM angular velocity in world frame. Shape: (N, 3).""" + return self.root_com_vel_w[:, 3:6] + + # -- Body state properties (link frame) -- + + @property + def body_link_pose_w(self) -> torch.Tensor: + """Body link poses in world frame. Shape: (N, num_bodies, 7).""" + if self._body_link_pose_w is None: + pose = torch.zeros(self._num_instances, self._num_bodies, 7, device=self._device) + pose[..., 3:7] = self._identity_quat(self._num_instances, self._num_bodies) + return pose + return self._body_link_pose_w + + @property + def body_link_vel_w(self) -> torch.Tensor: + """Body link velocities in world frame. Shape: (N, num_bodies, 6).""" + if self._body_link_vel_w is None: + return torch.zeros(self._num_instances, self._num_bodies, 6, device=self._device) + return self._body_link_vel_w + + @property + def body_link_state_w(self) -> torch.Tensor: + """Body link states in world frame. Shape: (N, num_bodies, 13).""" + return torch.cat([self.body_link_pose_w, self.body_link_vel_w], dim=-1) + + # Sliced properties + @property + def body_link_pos_w(self) -> torch.Tensor: + """Body link positions in world frame. Shape: (N, num_bodies, 3).""" + return self.body_link_pose_w[..., :3] + + @property + def body_link_quat_w(self) -> torch.Tensor: + """Body link orientations in world frame. Shape: (N, num_bodies, 4).""" + return self.body_link_pose_w[..., 3:7] + + @property + def body_link_lin_vel_w(self) -> torch.Tensor: + """Body link linear velocities in world frame. Shape: (N, num_bodies, 3).""" + return self.body_link_vel_w[..., :3] + + @property + def body_link_ang_vel_w(self) -> torch.Tensor: + """Body link angular velocities in world frame. Shape: (N, num_bodies, 3).""" + return self.body_link_vel_w[..., 3:6] + + # -- Body state properties (CoM frame) -- + + @property + def body_com_pose_w(self) -> torch.Tensor: + """Body CoM poses in world frame. Shape: (N, num_bodies, 7).""" + if self._body_com_pose_w is None: + return self.body_link_pose_w.clone() + return self._body_com_pose_w + + @property + def body_com_vel_w(self) -> torch.Tensor: + """Body CoM velocities in world frame. Shape: (N, num_bodies, 6).""" + if self._body_com_vel_w is None: + return self.body_link_vel_w.clone() + return self._body_com_vel_w + + @property + def body_com_state_w(self) -> torch.Tensor: + """Body CoM states in world frame. Shape: (N, num_bodies, 13).""" + return torch.cat([self.body_com_pose_w, self.body_com_vel_w], dim=-1) + + @property + def body_com_acc_w(self) -> torch.Tensor: + """Body CoM accelerations in world frame. Shape: (N, num_bodies, 6).""" + if self._body_com_acc_w is None: + return torch.zeros(self._num_instances, self._num_bodies, 6, device=self._device) + return self._body_com_acc_w + + @property + def body_com_pose_b(self) -> torch.Tensor: + """Body CoM poses in body frame. Shape: (N, 1, 7).""" + if self._body_com_pose_b is None: + pose = torch.zeros(self._num_instances, 1, 7, device=self._device) + pose[..., 3:7] = self._identity_quat(self._num_instances, 1) + return pose + return self._body_com_pose_b + + # Sliced properties + @property + def body_com_pos_w(self) -> torch.Tensor: + """Body CoM positions in world frame. Shape: (N, num_bodies, 3).""" + return self.body_com_pose_w[..., :3] + + @property + def body_com_quat_w(self) -> torch.Tensor: + """Body CoM orientations in world frame. Shape: (N, num_bodies, 4).""" + return self.body_com_pose_w[..., 3:7] + + @property + def body_com_lin_vel_w(self) -> torch.Tensor: + """Body CoM linear velocities in world frame. Shape: (N, num_bodies, 3).""" + return self.body_com_vel_w[..., :3] + + @property + def body_com_ang_vel_w(self) -> torch.Tensor: + """Body CoM angular velocities in world frame. Shape: (N, num_bodies, 3).""" + return self.body_com_vel_w[..., 3:6] + + @property + def body_com_lin_acc_w(self) -> torch.Tensor: + """Body CoM linear accelerations in world frame. Shape: (N, num_bodies, 3).""" + return self.body_com_acc_w[..., :3] + + @property + def body_com_ang_acc_w(self) -> torch.Tensor: + """Body CoM angular accelerations in world frame. Shape: (N, num_bodies, 3).""" + return self.body_com_acc_w[..., 3:6] + + @property + def body_com_pos_b(self) -> torch.Tensor: + """Body CoM positions in body frame. Shape: (N, num_bodies, 3).""" + return self.body_com_pose_b[..., :3] + + @property + def body_com_quat_b(self) -> torch.Tensor: + """Body CoM orientations in body frame. Shape: (N, num_bodies, 4).""" + return self.body_com_pose_b[..., 3:7] + + # -- Body properties -- + + @property + def body_mass(self) -> torch.Tensor: + """Body masses. Shape: (N, num_bodies).""" + if self._body_mass is None: + return torch.ones(self._num_instances, self._num_bodies, device=self._device) + return self._body_mass + + @property + def body_inertia(self) -> torch.Tensor: + """Body inertias. Shape: (N, num_bodies, 3, 3).""" + if self._body_inertia is None: + inertia = torch.zeros(self._num_instances, self._num_bodies, 3, 3, device=self._device) + inertia[..., 0, 0] = 1.0 + inertia[..., 1, 1] = 1.0 + inertia[..., 2, 2] = 1.0 + return inertia + return self._body_inertia + + @property + def body_incoming_joint_wrench_b(self) -> torch.Tensor: + """Body incoming joint wrenches. Shape: (N, num_bodies, 6).""" + if self._body_incoming_joint_wrench_b is None: + return torch.zeros(self._num_instances, self._num_bodies, 6, device=self._device) + return self._body_incoming_joint_wrench_b + + # -- Derived properties -- + + @property + def projected_gravity_b(self) -> torch.Tensor: + """Gravity projection on base. Shape: (N, 3).""" + # Default gravity pointing down + gravity = torch.zeros(self._num_instances, 3, device=self._device) + gravity[:, 2] = -1.0 + return gravity + + @property + def heading_w(self) -> torch.Tensor: + """Yaw heading in world frame. Shape: (N,).""" + return torch.zeros(self._num_instances, device=self._device) + + @property + def root_link_lin_vel_b(self) -> torch.Tensor: + """Root link linear velocity in body frame. Shape: (N, 3).""" + return self.root_link_lin_vel_w.clone() + + @property + def root_link_ang_vel_b(self) -> torch.Tensor: + """Root link angular velocity in body frame. Shape: (N, 3).""" + return self.root_link_ang_vel_w.clone() + + @property + def root_com_lin_vel_b(self) -> torch.Tensor: + """Root CoM linear velocity in body frame. Shape: (N, 3).""" + return self.root_com_lin_vel_w.clone() + + @property + def root_com_ang_vel_b(self) -> torch.Tensor: + """Root CoM angular velocity in body frame. Shape: (N, 3).""" + return self.root_com_ang_vel_w.clone() + + # -- Convenience aliases for root state (without _link_ or _com_ prefix) -- + + @property + def root_pos_w(self) -> torch.Tensor: + """Root position (alias for root_link_pos_w). Shape: (N, 3).""" + return self.root_link_pos_w + + @property + def root_quat_w(self) -> torch.Tensor: + """Root orientation (alias for root_link_quat_w). Shape: (N, 4).""" + return self.root_link_quat_w + + @property + def root_pose_w(self) -> torch.Tensor: + """Root pose (alias for root_link_pose_w). Shape: (N, 7).""" + return self.root_link_pose_w + + @property + def root_vel_w(self) -> torch.Tensor: + """Root velocity (alias for root_com_vel_w). Shape: (N, 6).""" + return self.root_com_vel_w + + @property + def root_lin_vel_w(self) -> torch.Tensor: + """Root linear velocity (alias for root_com_lin_vel_w). Shape: (N, 3).""" + return self.root_com_lin_vel_w + + @property + def root_ang_vel_w(self) -> torch.Tensor: + """Root angular velocity (alias for root_com_ang_vel_w). Shape: (N, 3).""" + return self.root_com_ang_vel_w + + @property + def root_lin_vel_b(self) -> torch.Tensor: + """Root linear velocity in body frame (alias for root_com_lin_vel_b). Shape: (N, 3).""" + return self.root_com_lin_vel_b + + @property + def root_ang_vel_b(self) -> torch.Tensor: + """Root angular velocity in body frame (alias for root_com_ang_vel_b). Shape: (N, 3).""" + return self.root_com_ang_vel_b + + # -- Convenience aliases for body state (without _link_ or _com_ prefix) -- + + @property + def body_pos_w(self) -> torch.Tensor: + """Body positions (alias for body_link_pos_w). Shape: (N, num_bodies, 3).""" + return self.body_link_pos_w + + @property + def body_quat_w(self) -> torch.Tensor: + """Body orientations (alias for body_link_quat_w). Shape: (N, num_bodies, 4).""" + return self.body_link_quat_w + + @property + def body_pose_w(self) -> torch.Tensor: + """Body poses (alias for body_link_pose_w). Shape: (N, num_bodies, 7).""" + return self.body_link_pose_w + + @property + def body_vel_w(self) -> torch.Tensor: + """Body velocities (alias for body_com_vel_w). Shape: (N, num_bodies, 6).""" + return self.body_com_vel_w + + @property + def body_state_w(self) -> torch.Tensor: + """Body states (alias for body_com_state_w). Shape: (N, num_bodies, 13).""" + return self.body_com_state_w + + @property + def body_lin_vel_w(self) -> torch.Tensor: + """Body linear velocities (alias for body_com_lin_vel_w). Shape: (N, num_bodies, 3).""" + return self.body_com_lin_vel_w + + @property + def body_ang_vel_w(self) -> torch.Tensor: + """Body angular velocities (alias for body_com_ang_vel_w). Shape: (N, num_bodies, 3).""" + return self.body_com_ang_vel_w + + @property + def body_acc_w(self) -> torch.Tensor: + """Body accelerations (alias for body_com_acc_w). Shape: (N, num_bodies, 6).""" + return self.body_com_acc_w + + @property + def body_lin_acc_w(self) -> torch.Tensor: + """Body linear accelerations (alias for body_com_lin_acc_w). Shape: (N, num_bodies, 3).""" + return self.body_com_lin_acc_w + + @property + def body_ang_acc_w(self) -> torch.Tensor: + """Body angular accelerations (alias for body_com_ang_acc_w). Shape: (N, num_bodies, 3).""" + return self.body_com_ang_acc_w + + # -- CoM in body frame (root body only) -- + + @property + def com_pos_b(self) -> torch.Tensor: + """Root CoM position in body frame. Shape: (N, 3).""" + return self.body_com_pose_b[:, 0, :3] + + @property + def com_quat_b(self) -> torch.Tensor: + """Root CoM orientation in body frame. Shape: (N, 4).""" + return self.body_com_pose_b[:, 0, 3:7] + + # -- Joint property aliases -- + + @property + def joint_limits(self) -> torch.Tensor: + """Joint position limits (alias for joint_pos_limits). Shape: (N, num_joints, 2).""" + return self.joint_pos_limits + + @property + def joint_velocity_limits(self) -> torch.Tensor: + """Joint velocity limits (alias for joint_vel_limits). Shape: (N, num_joints).""" + return self.joint_vel_limits + + @property + def joint_friction(self) -> torch.Tensor: + """Joint friction (alias for joint_friction_coeff). Shape: (N, num_joints).""" + return self.joint_friction_coeff + + # -- Fixed tendon alias -- + + @property + def fixed_tendon_limit(self) -> torch.Tensor: + """Fixed tendon limit (alias for fixed_tendon_pos_limits). Shape: (N, num_fixed_tendons, 2).""" + return self.fixed_tendon_pos_limits + + # -- Fixed tendon properties -- + + @property + def fixed_tendon_stiffness(self) -> torch.Tensor: + """Fixed tendon stiffness. Shape: (N, num_fixed_tendons).""" + if self._num_fixed_tendons == 0: + return torch.zeros(self._num_instances, 0, device=self._device) + if self._fixed_tendon_stiffness is None: + return torch.zeros(self._num_instances, self._num_fixed_tendons, device=self._device) + return self._fixed_tendon_stiffness + + @property + def fixed_tendon_damping(self) -> torch.Tensor: + """Fixed tendon damping. Shape: (N, num_fixed_tendons).""" + if self._num_fixed_tendons == 0: + return torch.zeros(self._num_instances, 0, device=self._device) + if self._fixed_tendon_damping is None: + return torch.zeros(self._num_instances, self._num_fixed_tendons, device=self._device) + return self._fixed_tendon_damping + + @property + def fixed_tendon_limit_stiffness(self) -> torch.Tensor: + """Fixed tendon limit stiffness. Shape: (N, num_fixed_tendons).""" + if self._num_fixed_tendons == 0: + return torch.zeros(self._num_instances, 0, device=self._device) + if self._fixed_tendon_limit_stiffness is None: + return torch.zeros(self._num_instances, self._num_fixed_tendons, device=self._device) + return self._fixed_tendon_limit_stiffness + + @property + def fixed_tendon_rest_length(self) -> torch.Tensor: + """Fixed tendon rest length. Shape: (N, num_fixed_tendons).""" + if self._num_fixed_tendons == 0: + return torch.zeros(self._num_instances, 0, device=self._device) + if self._fixed_tendon_rest_length is None: + return torch.zeros(self._num_instances, self._num_fixed_tendons, device=self._device) + return self._fixed_tendon_rest_length + + @property + def fixed_tendon_offset(self) -> torch.Tensor: + """Fixed tendon offset. Shape: (N, num_fixed_tendons).""" + if self._num_fixed_tendons == 0: + return torch.zeros(self._num_instances, 0, device=self._device) + if self._fixed_tendon_offset is None: + return torch.zeros(self._num_instances, self._num_fixed_tendons, device=self._device) + return self._fixed_tendon_offset + + @property + def fixed_tendon_pos_limits(self) -> torch.Tensor: + """Fixed tendon position limits. Shape: (N, num_fixed_tendons, 2).""" + if self._num_fixed_tendons == 0: + return torch.zeros(self._num_instances, 0, 2, device=self._device) + if self._fixed_tendon_pos_limits is None: + limits = torch.zeros(self._num_instances, self._num_fixed_tendons, 2, device=self._device) + limits[..., 0] = -float("inf") + limits[..., 1] = float("inf") + return limits + return self._fixed_tendon_pos_limits + + # -- Spatial tendon properties -- + + @property + def spatial_tendon_stiffness(self) -> torch.Tensor: + """Spatial tendon stiffness. Shape: (N, num_spatial_tendons).""" + if self._num_spatial_tendons == 0: + return torch.zeros(self._num_instances, 0, device=self._device) + if self._spatial_tendon_stiffness is None: + return torch.zeros(self._num_instances, self._num_spatial_tendons, device=self._device) + return self._spatial_tendon_stiffness + + @property + def spatial_tendon_damping(self) -> torch.Tensor: + """Spatial tendon damping. Shape: (N, num_spatial_tendons).""" + if self._num_spatial_tendons == 0: + return torch.zeros(self._num_instances, 0, device=self._device) + if self._spatial_tendon_damping is None: + return torch.zeros(self._num_instances, self._num_spatial_tendons, device=self._device) + return self._spatial_tendon_damping + + @property + def spatial_tendon_limit_stiffness(self) -> torch.Tensor: + """Spatial tendon limit stiffness. Shape: (N, num_spatial_tendons).""" + if self._num_spatial_tendons == 0: + return torch.zeros(self._num_instances, 0, device=self._device) + if self._spatial_tendon_limit_stiffness is None: + return torch.zeros(self._num_instances, self._num_spatial_tendons, device=self._device) + return self._spatial_tendon_limit_stiffness + + @property + def spatial_tendon_offset(self) -> torch.Tensor: + """Spatial tendon offset. Shape: (N, num_spatial_tendons).""" + if self._num_spatial_tendons == 0: + return torch.zeros(self._num_instances, 0, device=self._device) + if self._spatial_tendon_offset is None: + return torch.zeros(self._num_instances, self._num_spatial_tendons, device=self._device) + return self._spatial_tendon_offset + + # -- Update method -- + + def update(self, dt: float) -> None: + """Update data (no-op for mock).""" + pass + + # -- Setters -- + + def set_default_root_pose(self, value: torch.Tensor) -> None: + self._default_root_pose = value.to(self._device) + + def set_default_root_vel(self, value: torch.Tensor) -> None: + self._default_root_vel = value.to(self._device) + + def set_default_joint_pos(self, value: torch.Tensor) -> None: + self._default_joint_pos = value.to(self._device) + + def set_default_joint_vel(self, value: torch.Tensor) -> None: + self._default_joint_vel = value.to(self._device) + + def set_joint_pos_target(self, value: torch.Tensor) -> None: + self._joint_pos_target = value.to(self._device) + + def set_joint_vel_target(self, value: torch.Tensor) -> None: + self._joint_vel_target = value.to(self._device) + + def set_joint_effort_target(self, value: torch.Tensor) -> None: + self._joint_effort_target = value.to(self._device) + + def set_computed_torque(self, value: torch.Tensor) -> None: + self._computed_torque = value.to(self._device) + + def set_applied_torque(self, value: torch.Tensor) -> None: + self._applied_torque = value.to(self._device) + + def set_joint_stiffness(self, value: torch.Tensor) -> None: + self._joint_stiffness = value.to(self._device) + + def set_joint_damping(self, value: torch.Tensor) -> None: + self._joint_damping = value.to(self._device) + + def set_joint_armature(self, value: torch.Tensor) -> None: + self._joint_armature = value.to(self._device) + + def set_joint_friction_coeff(self, value: torch.Tensor) -> None: + self._joint_friction_coeff = value.to(self._device) + + def set_joint_pos_limits(self, value: torch.Tensor) -> None: + self._joint_pos_limits = value.to(self._device) + + def set_joint_vel_limits(self, value: torch.Tensor) -> None: + self._joint_vel_limits = value.to(self._device) + + def set_joint_effort_limits(self, value: torch.Tensor) -> None: + self._joint_effort_limits = value.to(self._device) + + def set_soft_joint_pos_limits(self, value: torch.Tensor) -> None: + self._soft_joint_pos_limits = value.to(self._device) + + def set_soft_joint_vel_limits(self, value: torch.Tensor) -> None: + self._soft_joint_vel_limits = value.to(self._device) + + def set_gear_ratio(self, value: torch.Tensor) -> None: + self._gear_ratio = value.to(self._device) + + def set_joint_pos(self, value: torch.Tensor) -> None: + self._joint_pos = value.to(self._device) + + def set_joint_vel(self, value: torch.Tensor) -> None: + self._joint_vel = value.to(self._device) + + def set_joint_acc(self, value: torch.Tensor) -> None: + self._joint_acc = value.to(self._device) + + def set_root_link_pose_w(self, value: torch.Tensor) -> None: + self._root_link_pose_w = value.to(self._device) + + def set_root_link_vel_w(self, value: torch.Tensor) -> None: + self._root_link_vel_w = value.to(self._device) + + def set_root_com_pose_w(self, value: torch.Tensor) -> None: + self._root_com_pose_w = value.to(self._device) + + def set_root_com_vel_w(self, value: torch.Tensor) -> None: + self._root_com_vel_w = value.to(self._device) + + def set_body_link_pose_w(self, value: torch.Tensor) -> None: + self._body_link_pose_w = value.to(self._device) + + def set_body_link_vel_w(self, value: torch.Tensor) -> None: + self._body_link_vel_w = value.to(self._device) + + def set_body_com_pose_w(self, value: torch.Tensor) -> None: + self._body_com_pose_w = value.to(self._device) + + def set_body_com_vel_w(self, value: torch.Tensor) -> None: + self._body_com_vel_w = value.to(self._device) + + def set_body_com_acc_w(self, value: torch.Tensor) -> None: + self._body_com_acc_w = value.to(self._device) + + def set_body_com_pose_b(self, value: torch.Tensor) -> None: + self._body_com_pose_b = value.to(self._device) + + def set_body_mass(self, value: torch.Tensor) -> None: + self._body_mass = value.to(self._device) + + def set_body_inertia(self, value: torch.Tensor) -> None: + self._body_inertia = value.to(self._device) + + def set_body_incoming_joint_wrench_b(self, value: torch.Tensor) -> None: + self._body_incoming_joint_wrench_b = value.to(self._device) + + def set_fixed_tendon_stiffness(self, value: torch.Tensor) -> None: + self._fixed_tendon_stiffness = value.to(self._device) + + def set_fixed_tendon_damping(self, value: torch.Tensor) -> None: + self._fixed_tendon_damping = value.to(self._device) + + def set_fixed_tendon_limit_stiffness(self, value: torch.Tensor) -> None: + self._fixed_tendon_limit_stiffness = value.to(self._device) + + def set_fixed_tendon_rest_length(self, value: torch.Tensor) -> None: + self._fixed_tendon_rest_length = value.to(self._device) + + def set_fixed_tendon_offset(self, value: torch.Tensor) -> None: + self._fixed_tendon_offset = value.to(self._device) + + def set_fixed_tendon_pos_limits(self, value: torch.Tensor) -> None: + self._fixed_tendon_pos_limits = value.to(self._device) + + def set_spatial_tendon_stiffness(self, value: torch.Tensor) -> None: + self._spatial_tendon_stiffness = value.to(self._device) + + def set_spatial_tendon_damping(self, value: torch.Tensor) -> None: + self._spatial_tendon_damping = value.to(self._device) + + def set_spatial_tendon_limit_stiffness(self, value: torch.Tensor) -> None: + self._spatial_tendon_limit_stiffness = value.to(self._device) + + def set_spatial_tendon_offset(self, value: torch.Tensor) -> None: + self._spatial_tendon_offset = value.to(self._device) + + def set_mock_data(self, **kwargs) -> None: + """Bulk setter for mock data. + + Accepts any property name as a keyword argument with a tensor value. + """ + for key, value in kwargs.items(): + setter = getattr(self, f"set_{key}", None) + if setter is not None: + setter(value) + else: + raise ValueError(f"Unknown property: {key}") + + +class MockArticulation: + """Mock articulation asset for testing without Isaac Sim. + + This class mimics the interface of BaseArticulation for testing purposes. + It provides the same properties and methods but without simulation dependencies. + """ + + def __init__( + self, + num_instances: int, + num_joints: int, + num_bodies: int, + joint_names: list[str] | None = None, + body_names: list[str] | None = None, + is_fixed_base: bool = False, + num_fixed_tendons: int = 0, + num_spatial_tendons: int = 0, + fixed_tendon_names: list[str] | None = None, + spatial_tendon_names: list[str] | None = None, + device: str = "cpu", + ): + """Initialize mock articulation. + + Args: + num_instances: Number of articulation instances. + num_joints: Number of joints. + num_bodies: Number of bodies. + joint_names: Names of joints. + body_names: Names of bodies. + is_fixed_base: Whether the articulation has a fixed base. + num_fixed_tendons: Number of fixed tendons. + num_spatial_tendons: Number of spatial tendons. + fixed_tendon_names: Names of fixed tendons. + spatial_tendon_names: Names of spatial tendons. + device: Device for tensor allocation. + """ + self._num_instances = num_instances + self._num_joints = num_joints + self._num_bodies = num_bodies + self._is_fixed_base = is_fixed_base + self._num_fixed_tendons = num_fixed_tendons + self._num_spatial_tendons = num_spatial_tendons + self._device = device + + self._joint_names = joint_names or [f"joint_{i}" for i in range(num_joints)] + self._body_names = body_names or [f"body_{i}" for i in range(num_bodies)] + self._fixed_tendon_names = fixed_tendon_names or [f"fixed_tendon_{i}" for i in range(num_fixed_tendons)] + self._spatial_tendon_names = spatial_tendon_names or [f"spatial_tendon_{i}" for i in range(num_spatial_tendons)] + + self._data = MockArticulationData( + num_instances=num_instances, + num_joints=num_joints, + num_bodies=num_bodies, + joint_names=self._joint_names, + body_names=self._body_names, + num_fixed_tendons=num_fixed_tendons, + num_spatial_tendons=num_spatial_tendons, + fixed_tendon_names=self._fixed_tendon_names, + spatial_tendon_names=self._spatial_tendon_names, + device=device, + ) + + # Actuators (empty dict for mock) + self.actuators: dict = {} + + # -- Properties -- + + @property + def data(self) -> MockArticulationData: + """Data container for the articulation.""" + return self._data + + @property + def num_instances(self) -> int: + """Number of articulation instances.""" + return self._num_instances + + @property + def device(self) -> str: + """Device for tensor allocation.""" + return self._device + + @property + def is_fixed_base(self) -> bool: + """Whether the articulation has a fixed base.""" + return self._is_fixed_base + + @property + def num_joints(self) -> int: + """Number of joints.""" + return self._num_joints + + @property + def num_bodies(self) -> int: + """Number of bodies.""" + return self._num_bodies + + @property + def num_fixed_tendons(self) -> int: + """Number of fixed tendons.""" + return self._num_fixed_tendons + + @property + def num_spatial_tendons(self) -> int: + """Number of spatial tendons.""" + return self._num_spatial_tendons + + @property + def joint_names(self) -> list[str]: + """Ordered joint names.""" + return self._joint_names + + @property + def body_names(self) -> list[str]: + """Ordered body names.""" + return self._body_names + + @property + def fixed_tendon_names(self) -> list[str]: + """Ordered fixed tendon names.""" + return self._fixed_tendon_names + + @property + def spatial_tendon_names(self) -> list[str]: + """Ordered spatial tendon names.""" + return self._spatial_tendon_names + + @property + def root_view(self) -> None: + """Returns None (no physics view in mock).""" + return None + + @property + def instantaneous_wrench_composer(self) -> None: + """Returns None (no wrench composer in mock).""" + return None + + @property + def permanent_wrench_composer(self) -> None: + """Returns None (no wrench composer in mock).""" + return None + + # -- Core methods -- + + def reset(self, env_ids: Sequence[int] | None = None) -> None: + """Reset articulation state for specified environments.""" + pass + + def write_data_to_sim(self) -> None: + """Write data to simulation (no-op for mock).""" + pass + + def update(self, dt: float) -> None: + """Update articulation data.""" + self._data.update(dt) + + # -- Finder methods -- + + def find_bodies(self, name_keys: str | Sequence[str], preserve_order: bool = False) -> tuple[list[int], list[str]]: + """Find bodies by name regex patterns.""" + return self._find_by_regex(self._body_names, name_keys, preserve_order) + + def find_joints( + self, + name_keys: str | Sequence[str], + joint_subset: list[int] | None = None, + preserve_order: bool = False, + ) -> tuple[list[int], list[str]]: + """Find joints by name regex patterns.""" + names = self._joint_names + if joint_subset is not None: + names = [names[i] for i in joint_subset] + indices, matched_names = self._find_by_regex(names, name_keys, preserve_order) + if joint_subset is not None: + indices = [joint_subset[i] for i in indices] + return indices, matched_names + + def find_fixed_tendons( + self, + name_keys: str | Sequence[str], + tendon_subsets: list[int] | None = None, + preserve_order: bool = False, + ) -> tuple[list[int], list[str]]: + """Find fixed tendons by name regex patterns.""" + names = self._fixed_tendon_names + if tendon_subsets is not None: + names = [names[i] for i in tendon_subsets] + indices, matched_names = self._find_by_regex(names, name_keys, preserve_order) + if tendon_subsets is not None: + indices = [tendon_subsets[i] for i in indices] + return indices, matched_names + + def find_spatial_tendons( + self, + name_keys: str | Sequence[str], + tendon_subsets: list[int] | None = None, + preserve_order: bool = False, + ) -> tuple[list[int], list[str]]: + """Find spatial tendons by name regex patterns.""" + names = self._spatial_tendon_names + if tendon_subsets is not None: + names = [names[i] for i in tendon_subsets] + indices, matched_names = self._find_by_regex(names, name_keys, preserve_order) + if tendon_subsets is not None: + indices = [tendon_subsets[i] for i in indices] + return indices, matched_names + + def _find_by_regex( + self, names: list[str], name_keys: str | Sequence[str], preserve_order: bool + ) -> tuple[list[int], list[str]]: + """Find items by regex patterns.""" + if isinstance(name_keys, str): + name_keys = [name_keys] + + matched_indices = [] + matched_names = [] + + if preserve_order: + for key in name_keys: + pattern = re.compile(key) + for i, name in enumerate(names): + if pattern.fullmatch(name) and i not in matched_indices: + matched_indices.append(i) + matched_names.append(name) + else: + for i, name in enumerate(names): + for key in name_keys: + pattern = re.compile(key) + if pattern.fullmatch(name): + matched_indices.append(i) + matched_names.append(name) + break + + return matched_indices, matched_names + + # -- State writer methods (no-op for mock) -- + + def write_root_state_to_sim( + self, + root_state: torch.Tensor, + env_ids: Sequence[int] | None = None, + ) -> None: + """Write root state to simulation.""" + pass + + def write_root_com_state_to_sim( + self, + root_state: torch.Tensor, + env_ids: Sequence[int] | None = None, + ) -> None: + """Write root CoM state to simulation.""" + pass + + def write_root_link_state_to_sim( + self, + root_state: torch.Tensor, + env_ids: Sequence[int] | None = None, + ) -> None: + """Write root link state to simulation.""" + pass + + def write_root_pose_to_sim( + self, + root_pose: torch.Tensor, + env_ids: Sequence[int] | None = None, + ) -> None: + """Write root pose to simulation.""" + pass + + def write_root_link_pose_to_sim( + self, + root_pose: torch.Tensor, + env_ids: Sequence[int] | None = None, + ) -> None: + """Write root link pose to simulation.""" + pass + + def write_root_com_pose_to_sim( + self, + root_pose: torch.Tensor, + env_ids: Sequence[int] | None = None, + ) -> None: + """Write root CoM pose to simulation.""" + pass + + def write_root_velocity_to_sim( + self, + root_velocity: torch.Tensor, + env_ids: Sequence[int] | None = None, + ) -> None: + """Write root velocity to simulation.""" + pass + + def write_root_com_velocity_to_sim( + self, + root_velocity: torch.Tensor, + env_ids: Sequence[int] | None = None, + ) -> None: + """Write root CoM velocity to simulation.""" + pass + + def write_root_link_velocity_to_sim( + self, + root_velocity: torch.Tensor, + env_ids: Sequence[int] | None = None, + ) -> None: + """Write root link velocity to simulation.""" + pass + + def write_joint_state_to_sim( + self, + position: torch.Tensor, + velocity: torch.Tensor, + joint_ids: Sequence[int] | slice | None = None, + env_ids: Sequence[int] | None = None, + ) -> None: + """Write joint state to simulation.""" + pass + + def write_joint_position_to_sim( + self, + position: torch.Tensor, + joint_ids: Sequence[int] | slice | None = None, + env_ids: Sequence[int] | None = None, + ) -> None: + """Write joint positions to simulation.""" + pass + + def write_joint_velocity_to_sim( + self, + velocity: torch.Tensor, + joint_ids: Sequence[int] | slice | None = None, + env_ids: Sequence[int] | None = None, + ) -> None: + """Write joint velocities to simulation.""" + pass + + def write_joint_stiffness_to_sim( + self, + stiffness: torch.Tensor | float, + joint_ids: Sequence[int] | slice | None = None, + env_ids: Sequence[int] | None = None, + ) -> None: + """Write joint stiffness to simulation.""" + pass + + def write_joint_damping_to_sim( + self, + damping: torch.Tensor | float, + joint_ids: Sequence[int] | slice | None = None, + env_ids: Sequence[int] | None = None, + ) -> None: + """Write joint damping to simulation.""" + pass + + def write_joint_position_limit_to_sim( + self, + limits: torch.Tensor | float, + joint_ids: Sequence[int] | slice | None = None, + env_ids: Sequence[int] | None = None, + warn_limit_violation: bool = True, + ) -> None: + """Write joint position limits to simulation.""" + pass + + def write_joint_velocity_limit_to_sim( + self, + limits: torch.Tensor | float, + joint_ids: Sequence[int] | slice | None = None, + env_ids: Sequence[int] | None = None, + ) -> None: + """Write joint velocity limits to simulation.""" + pass + + def write_joint_effort_limit_to_sim( + self, + limits: torch.Tensor | float, + joint_ids: Sequence[int] | slice | None = None, + env_ids: Sequence[int] | None = None, + ) -> None: + """Write joint effort limits to simulation.""" + pass + + def write_joint_armature_to_sim( + self, + armature: torch.Tensor | float, + joint_ids: Sequence[int] | slice | None = None, + env_ids: Sequence[int] | None = None, + ) -> None: + """Write joint armature to simulation.""" + pass + + def write_joint_friction_coefficient_to_sim( + self, + coeff: torch.Tensor | float, + joint_ids: Sequence[int] | slice | None = None, + env_ids: Sequence[int] | None = None, + ) -> None: + """Write joint friction coefficient to simulation.""" + pass + + def write_joint_friction_to_sim( + self, + joint_ids: Sequence[int] | slice | None = None, + env_ids: Sequence[int] | None = None, + ) -> None: + """Write joint friction to simulation.""" + pass + + def write_joint_limits_to_sim( + self, + joint_ids: Sequence[int] | slice | None = None, + env_ids: Sequence[int] | None = None, + ) -> None: + """Write joint limits to simulation.""" + pass + + # -- Setter methods -- + + def set_masses( + self, + masses: torch.Tensor, + body_ids: Sequence[int] | slice | None = None, + env_ids: Sequence[int] | None = None, + ) -> None: + """Set body masses.""" + pass + + def set_coms( + self, + coms: torch.Tensor, + body_ids: Sequence[int] | slice | None = None, + env_ids: Sequence[int] | None = None, + ) -> None: + """Set body centers of mass.""" + pass + + def set_inertias( + self, + inertias: torch.Tensor, + body_ids: Sequence[int] | slice | None = None, + env_ids: Sequence[int] | None = None, + ) -> None: + """Set body inertias.""" + pass + + def set_external_force_and_torque( + self, + forces: torch.Tensor, + torques: torch.Tensor, + positions: torch.Tensor | None = None, + body_ids: Sequence[int] | slice | None = None, + env_ids: Sequence[int] | None = None, + is_global: bool = True, + ) -> None: + """Set external forces and torques.""" + pass + + def set_joint_position_target( + self, + target: torch.Tensor, + joint_ids: Sequence[int] | slice | None = None, + env_ids: Sequence[int] | None = None, + ) -> None: + """Set joint position targets.""" + if env_ids is None: + env_ids = slice(None) + if joint_ids is None: + joint_ids = slice(None) + if self._data._joint_pos_target is None: + self._data._joint_pos_target = torch.zeros(self._num_instances, self._num_joints, device=self._device) + self._data._joint_pos_target[env_ids, joint_ids] = target.to(self._device) + + def set_joint_velocity_target( + self, + target: torch.Tensor, + joint_ids: Sequence[int] | slice | None = None, + env_ids: Sequence[int] | None = None, + ) -> None: + """Set joint velocity targets.""" + if env_ids is None: + env_ids = slice(None) + if joint_ids is None: + joint_ids = slice(None) + if self._data._joint_vel_target is None: + self._data._joint_vel_target = torch.zeros(self._num_instances, self._num_joints, device=self._device) + self._data._joint_vel_target[env_ids, joint_ids] = target.to(self._device) + + def set_joint_effort_target( + self, + target: torch.Tensor, + joint_ids: Sequence[int] | slice | None = None, + env_ids: Sequence[int] | None = None, + ) -> None: + """Set joint effort targets.""" + if env_ids is None: + env_ids = slice(None) + if joint_ids is None: + joint_ids = slice(None) + if self._data._joint_effort_target is None: + self._data._joint_effort_target = torch.zeros(self._num_instances, self._num_joints, device=self._device) + self._data._joint_effort_target[env_ids, joint_ids] = target.to(self._device) + + # -- Tendon methods (fixed) -- + + def set_fixed_tendon_stiffness( + self, + stiffness: torch.Tensor, + tendon_ids: Sequence[int] | slice | None = None, + env_ids: Sequence[int] | None = None, + ) -> None: + """Set fixed tendon stiffness.""" + pass + + def set_fixed_tendon_damping( + self, + damping: torch.Tensor, + tendon_ids: Sequence[int] | slice | None = None, + env_ids: Sequence[int] | None = None, + ) -> None: + """Set fixed tendon damping.""" + pass + + def set_fixed_tendon_limit_stiffness( + self, + limit_stiffness: torch.Tensor, + tendon_ids: Sequence[int] | slice | None = None, + env_ids: Sequence[int] | None = None, + ) -> None: + """Set fixed tendon limit stiffness.""" + pass + + def set_fixed_tendon_position_limit( + self, + limit: torch.Tensor, + tendon_ids: Sequence[int] | slice | None = None, + env_ids: Sequence[int] | None = None, + ) -> None: + """Set fixed tendon position limit.""" + pass + + def set_fixed_tendon_limit( + self, + limit: torch.Tensor, + tendon_ids: Sequence[int] | slice | None = None, + env_ids: Sequence[int] | None = None, + ) -> None: + """Set fixed tendon limit (alias for set_fixed_tendon_position_limit).""" + pass + + def set_fixed_tendon_rest_length( + self, + rest_length: torch.Tensor, + tendon_ids: Sequence[int] | slice | None = None, + env_ids: Sequence[int] | None = None, + ) -> None: + """Set fixed tendon rest length.""" + pass + + def set_fixed_tendon_offset( + self, + offset: torch.Tensor, + tendon_ids: Sequence[int] | slice | None = None, + env_ids: Sequence[int] | None = None, + ) -> None: + """Set fixed tendon offset.""" + pass + + def write_fixed_tendon_properties_to_sim( + self, + tendon_ids: Sequence[int] | slice | None = None, + env_ids: Sequence[int] | None = None, + ) -> None: + """Write fixed tendon properties to simulation.""" + pass + + # -- Tendon methods (spatial) -- + + def set_spatial_tendon_stiffness( + self, + stiffness: torch.Tensor, + tendon_ids: Sequence[int] | slice | None = None, + env_ids: Sequence[int] | None = None, + ) -> None: + """Set spatial tendon stiffness.""" + pass + + def set_spatial_tendon_damping( + self, + damping: torch.Tensor, + tendon_ids: Sequence[int] | slice | None = None, + env_ids: Sequence[int] | None = None, + ) -> None: + """Set spatial tendon damping.""" + pass + + def set_spatial_tendon_limit_stiffness( + self, + limit_stiffness: torch.Tensor, + tendon_ids: Sequence[int] | slice | None = None, + env_ids: Sequence[int] | None = None, + ) -> None: + """Set spatial tendon limit stiffness.""" + pass + + def set_spatial_tendon_offset( + self, + offset: torch.Tensor, + tendon_ids: Sequence[int] | slice | None = None, + env_ids: Sequence[int] | None = None, + ) -> None: + """Set spatial tendon offset.""" + pass + + def write_spatial_tendon_properties_to_sim( + self, + tendon_ids: Sequence[int] | slice | None = None, + env_ids: Sequence[int] | None = None, + ) -> None: + """Write spatial tendon properties to simulation.""" + pass diff --git a/source/isaaclab/isaaclab/test/mock_interfaces/assets/mock_rigid_object.py b/source/isaaclab/isaaclab/test/mock_interfaces/assets/mock_rigid_object.py new file mode 100644 index 00000000000..8e9354b8814 --- /dev/null +++ b/source/isaaclab/isaaclab/test/mock_interfaces/assets/mock_rigid_object.py @@ -0,0 +1,759 @@ +# 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 + +"""Mock rigid object asset for testing without Isaac Sim.""" + +from __future__ import annotations + +import re +from collections.abc import Sequence + +import torch + + +class MockRigidObjectData: + """Mock data container for rigid object asset. + + This class mimics the interface of BaseRigidObjectData for testing purposes. + All tensor properties return zero tensors with correct shapes if not explicitly set. + """ + + def __init__( + self, + num_instances: int, + body_names: list[str] | None = None, + device: str = "cpu", + ): + """Initialize mock rigid object data. + + Args: + num_instances: Number of rigid object instances. + body_names: Names of bodies (single body for rigid object). + device: Device for tensor allocation. + """ + self._num_instances = num_instances + self._num_bodies = 1 # Rigid object always has 1 body + self._device = device + + self.body_names = body_names or ["body"] + + # Default states + self._default_root_pose: torch.Tensor | None = None + self._default_root_vel: torch.Tensor | None = None + + # Root state (link frame) + self._root_link_pose_w: torch.Tensor | None = None + self._root_link_vel_w: torch.Tensor | None = None + + # Root state (CoM frame) + self._root_com_pose_w: torch.Tensor | None = None + self._root_com_vel_w: torch.Tensor | None = None + + # Body state (link frame) + self._body_link_pose_w: torch.Tensor | None = None + self._body_link_vel_w: torch.Tensor | None = None + + # Body state (CoM frame) + self._body_com_pose_w: torch.Tensor | None = None + self._body_com_vel_w: torch.Tensor | None = None + self._body_com_acc_w: torch.Tensor | None = None + self._body_com_pose_b: torch.Tensor | None = None + + # Body properties + self._body_mass: torch.Tensor | None = None + self._body_inertia: torch.Tensor | None = None + + @property + def device(self) -> str: + """Device for tensor allocation.""" + return self._device + + def _identity_quat(self, *shape: int) -> torch.Tensor: + """Create identity quaternion tensor (w, x, y, z) = (1, 0, 0, 0).""" + quat = torch.zeros(*shape, 4, device=self._device) + quat[..., 0] = 1.0 + return quat + + # -- Default state properties -- + + @property + def default_root_pose(self) -> torch.Tensor: + """Default root pose. Shape: (N, 7).""" + if self._default_root_pose is None: + pose = torch.zeros(self._num_instances, 7, device=self._device) + pose[:, 3:7] = self._identity_quat(self._num_instances) + return pose + return self._default_root_pose + + @property + def default_root_vel(self) -> torch.Tensor: + """Default root velocity. Shape: (N, 6).""" + if self._default_root_vel is None: + return torch.zeros(self._num_instances, 6, device=self._device) + return self._default_root_vel + + @property + def default_root_state(self) -> torch.Tensor: + """Default root state. Shape: (N, 13).""" + return torch.cat([self.default_root_pose, self.default_root_vel], dim=-1) + + # -- Root state properties (link frame) -- + + @property + def root_link_pose_w(self) -> torch.Tensor: + """Root link pose in world frame. Shape: (N, 7).""" + if self._root_link_pose_w is None: + pose = torch.zeros(self._num_instances, 7, device=self._device) + pose[:, 3:7] = self._identity_quat(self._num_instances) + return pose + return self._root_link_pose_w + + @property + def root_link_vel_w(self) -> torch.Tensor: + """Root link velocity in world frame. Shape: (N, 6).""" + if self._root_link_vel_w is None: + return torch.zeros(self._num_instances, 6, device=self._device) + return self._root_link_vel_w + + @property + def root_link_state_w(self) -> torch.Tensor: + """Root link state in world frame. Shape: (N, 13).""" + return torch.cat([self.root_link_pose_w, self.root_link_vel_w], dim=-1) + + @property + def root_link_pos_w(self) -> torch.Tensor: + """Root link position. Shape: (N, 3).""" + return self.root_link_pose_w[:, :3] + + @property + def root_link_quat_w(self) -> torch.Tensor: + """Root link orientation. Shape: (N, 4).""" + return self.root_link_pose_w[:, 3:7] + + @property + def root_link_lin_vel_w(self) -> torch.Tensor: + """Root link linear velocity. Shape: (N, 3).""" + return self.root_link_vel_w[:, :3] + + @property + def root_link_ang_vel_w(self) -> torch.Tensor: + """Root link angular velocity. Shape: (N, 3).""" + return self.root_link_vel_w[:, 3:6] + + # -- Root state properties (CoM frame) -- + + @property + def root_com_pose_w(self) -> torch.Tensor: + """Root CoM pose in world frame. Shape: (N, 7).""" + if self._root_com_pose_w is None: + return self.root_link_pose_w.clone() + return self._root_com_pose_w + + @property + def root_com_vel_w(self) -> torch.Tensor: + """Root CoM velocity in world frame. Shape: (N, 6).""" + if self._root_com_vel_w is None: + return self.root_link_vel_w.clone() + return self._root_com_vel_w + + @property + def root_com_state_w(self) -> torch.Tensor: + """Root CoM state in world frame. Shape: (N, 13).""" + return torch.cat([self.root_com_pose_w, self.root_com_vel_w], dim=-1) + + @property + def root_state_w(self) -> torch.Tensor: + """Root state (link pose + CoM velocity). Shape: (N, 13).""" + return torch.cat([self.root_link_pose_w, self.root_com_vel_w], dim=-1) + + @property + def root_com_pos_w(self) -> torch.Tensor: + """Root CoM position. Shape: (N, 3).""" + return self.root_com_pose_w[:, :3] + + @property + def root_com_quat_w(self) -> torch.Tensor: + """Root CoM orientation. Shape: (N, 4).""" + return self.root_com_pose_w[:, 3:7] + + @property + def root_com_lin_vel_w(self) -> torch.Tensor: + """Root CoM linear velocity. Shape: (N, 3).""" + return self.root_com_vel_w[:, :3] + + @property + def root_com_ang_vel_w(self) -> torch.Tensor: + """Root CoM angular velocity. Shape: (N, 3).""" + return self.root_com_vel_w[:, 3:6] + + # -- Body state properties (link frame) -- + + @property + def body_link_pose_w(self) -> torch.Tensor: + """Body link pose in world frame. Shape: (N, 1, 7).""" + if self._body_link_pose_w is None: + pose = torch.zeros(self._num_instances, 1, 7, device=self._device) + pose[..., 3:7] = self._identity_quat(self._num_instances, 1) + return pose + return self._body_link_pose_w + + @property + def body_link_vel_w(self) -> torch.Tensor: + """Body link velocity in world frame. Shape: (N, 1, 6).""" + if self._body_link_vel_w is None: + return torch.zeros(self._num_instances, 1, 6, device=self._device) + return self._body_link_vel_w + + @property + def body_link_state_w(self) -> torch.Tensor: + """Body link state in world frame. Shape: (N, 1, 13).""" + return torch.cat([self.body_link_pose_w, self.body_link_vel_w], dim=-1) + + @property + def body_link_pos_w(self) -> torch.Tensor: + """Body link position. Shape: (N, 1, 3).""" + return self.body_link_pose_w[..., :3] + + @property + def body_link_quat_w(self) -> torch.Tensor: + """Body link orientation. Shape: (N, 1, 4).""" + return self.body_link_pose_w[..., 3:7] + + @property + def body_link_lin_vel_w(self) -> torch.Tensor: + """Body link linear velocity. Shape: (N, 1, 3).""" + return self.body_link_vel_w[..., :3] + + @property + def body_link_ang_vel_w(self) -> torch.Tensor: + """Body link angular velocity. Shape: (N, 1, 3).""" + return self.body_link_vel_w[..., 3:6] + + # -- Body state properties (CoM frame) -- + + @property + def body_com_pose_w(self) -> torch.Tensor: + """Body CoM pose in world frame. Shape: (N, 1, 7).""" + if self._body_com_pose_w is None: + return self.body_link_pose_w.clone() + return self._body_com_pose_w + + @property + def body_com_vel_w(self) -> torch.Tensor: + """Body CoM velocity in world frame. Shape: (N, 1, 6).""" + if self._body_com_vel_w is None: + return self.body_link_vel_w.clone() + return self._body_com_vel_w + + @property + def body_com_state_w(self) -> torch.Tensor: + """Body CoM state in world frame. Shape: (N, 1, 13).""" + return torch.cat([self.body_com_pose_w, self.body_com_vel_w], dim=-1) + + @property + def body_com_acc_w(self) -> torch.Tensor: + """Body CoM acceleration in world frame. Shape: (N, 1, 6).""" + if self._body_com_acc_w is None: + return torch.zeros(self._num_instances, 1, 6, device=self._device) + return self._body_com_acc_w + + @property + def body_com_pose_b(self) -> torch.Tensor: + """Body CoM pose in body frame. Shape: (N, 1, 7).""" + if self._body_com_pose_b is None: + pose = torch.zeros(self._num_instances, 1, 7, device=self._device) + pose[..., 3:7] = self._identity_quat(self._num_instances, 1) + return pose + return self._body_com_pose_b + + @property + def body_com_pos_w(self) -> torch.Tensor: + """Body CoM position. Shape: (N, 1, 3).""" + return self.body_com_pose_w[..., :3] + + @property + def body_com_quat_w(self) -> torch.Tensor: + """Body CoM orientation. Shape: (N, 1, 4).""" + return self.body_com_pose_w[..., 3:7] + + @property + def body_com_lin_vel_w(self) -> torch.Tensor: + """Body CoM linear velocity. Shape: (N, 1, 3).""" + return self.body_com_vel_w[..., :3] + + @property + def body_com_ang_vel_w(self) -> torch.Tensor: + """Body CoM angular velocity. Shape: (N, 1, 3).""" + return self.body_com_vel_w[..., 3:6] + + @property + def body_com_lin_acc_w(self) -> torch.Tensor: + """Body CoM linear acceleration. Shape: (N, 1, 3).""" + return self.body_com_acc_w[..., :3] + + @property + def body_com_ang_acc_w(self) -> torch.Tensor: + """Body CoM angular acceleration. Shape: (N, 1, 3).""" + return self.body_com_acc_w[..., 3:6] + + @property + def body_com_pos_b(self) -> torch.Tensor: + """Body CoM position in body frame. Shape: (N, 1, 3).""" + return self.body_com_pose_b[..., :3] + + @property + def body_com_quat_b(self) -> torch.Tensor: + """Body CoM orientation in body frame. Shape: (N, 1, 4).""" + return self.body_com_pose_b[..., 3:7] + + # -- Body properties -- + + @property + def body_mass(self) -> torch.Tensor: + """Body mass. Shape: (N, 1, 1).""" + if self._body_mass is None: + return torch.ones(self._num_instances, 1, 1, device=self._device) + return self._body_mass + + @property + def body_inertia(self) -> torch.Tensor: + """Body inertia (flattened 3x3). Shape: (N, 1, 9).""" + if self._body_inertia is None: + inertia = torch.zeros(self._num_instances, 1, 9, device=self._device) + inertia[..., 0] = 1.0 # Ixx + inertia[..., 4] = 1.0 # Iyy + inertia[..., 8] = 1.0 # Izz + return inertia + return self._body_inertia + + # -- Derived properties -- + + @property + def projected_gravity_b(self) -> torch.Tensor: + """Gravity projection on body. Shape: (N, 3).""" + gravity = torch.zeros(self._num_instances, 3, device=self._device) + gravity[:, 2] = -1.0 + return gravity + + @property + def heading_w(self) -> torch.Tensor: + """Yaw heading in world frame. Shape: (N,).""" + return torch.zeros(self._num_instances, device=self._device) + + @property + def root_link_lin_vel_b(self) -> torch.Tensor: + """Root link linear velocity in body frame. Shape: (N, 3).""" + return self.root_link_lin_vel_w.clone() + + @property + def root_link_ang_vel_b(self) -> torch.Tensor: + """Root link angular velocity in body frame. Shape: (N, 3).""" + return self.root_link_ang_vel_w.clone() + + @property + def root_com_lin_vel_b(self) -> torch.Tensor: + """Root CoM linear velocity in body frame. Shape: (N, 3).""" + return self.root_com_lin_vel_w.clone() + + @property + def root_com_ang_vel_b(self) -> torch.Tensor: + """Root CoM angular velocity in body frame. Shape: (N, 3).""" + return self.root_com_ang_vel_w.clone() + + # -- Convenience aliases for root state (without _link_ or _com_ prefix) -- + + @property + def root_pos_w(self) -> torch.Tensor: + """Root position (alias for root_link_pos_w). Shape: (N, 3).""" + return self.root_link_pos_w + + @property + def root_quat_w(self) -> torch.Tensor: + """Root orientation (alias for root_link_quat_w). Shape: (N, 4).""" + return self.root_link_quat_w + + @property + def root_pose_w(self) -> torch.Tensor: + """Root pose (alias for root_link_pose_w). Shape: (N, 7).""" + return self.root_link_pose_w + + @property + def root_vel_w(self) -> torch.Tensor: + """Root velocity (alias for root_com_vel_w). Shape: (N, 6).""" + return self.root_com_vel_w + + @property + def root_lin_vel_w(self) -> torch.Tensor: + """Root linear velocity (alias for root_com_lin_vel_w). Shape: (N, 3).""" + return self.root_com_lin_vel_w + + @property + def root_ang_vel_w(self) -> torch.Tensor: + """Root angular velocity (alias for root_com_ang_vel_w). Shape: (N, 3).""" + return self.root_com_ang_vel_w + + @property + def root_lin_vel_b(self) -> torch.Tensor: + """Root linear velocity in body frame (alias for root_com_lin_vel_b). Shape: (N, 3).""" + return self.root_com_lin_vel_b + + @property + def root_ang_vel_b(self) -> torch.Tensor: + """Root angular velocity in body frame (alias for root_com_ang_vel_b). Shape: (N, 3).""" + return self.root_com_ang_vel_b + + # -- Convenience aliases for body state (without _link_ or _com_ prefix) -- + + @property + def body_pos_w(self) -> torch.Tensor: + """Body positions (alias for body_link_pos_w). Shape: (N, 1, 3).""" + return self.body_link_pos_w + + @property + def body_quat_w(self) -> torch.Tensor: + """Body orientations (alias for body_link_quat_w). Shape: (N, 1, 4).""" + return self.body_link_quat_w + + @property + def body_pose_w(self) -> torch.Tensor: + """Body poses (alias for body_link_pose_w). Shape: (N, 1, 7).""" + return self.body_link_pose_w + + @property + def body_vel_w(self) -> torch.Tensor: + """Body velocities (alias for body_com_vel_w). Shape: (N, 1, 6).""" + return self.body_com_vel_w + + @property + def body_state_w(self) -> torch.Tensor: + """Body states (alias for body_com_state_w). Shape: (N, 1, 13).""" + return self.body_com_state_w + + @property + def body_lin_vel_w(self) -> torch.Tensor: + """Body linear velocities (alias for body_com_lin_vel_w). Shape: (N, 1, 3).""" + return self.body_com_lin_vel_w + + @property + def body_ang_vel_w(self) -> torch.Tensor: + """Body angular velocities (alias for body_com_ang_vel_w). Shape: (N, 1, 3).""" + return self.body_com_ang_vel_w + + @property + def body_acc_w(self) -> torch.Tensor: + """Body accelerations (alias for body_com_acc_w). Shape: (N, 1, 6).""" + return self.body_com_acc_w + + @property + def body_lin_acc_w(self) -> torch.Tensor: + """Body linear accelerations (alias for body_com_lin_acc_w). Shape: (N, 1, 3).""" + return self.body_com_lin_acc_w + + @property + def body_ang_acc_w(self) -> torch.Tensor: + """Body angular accelerations (alias for body_com_ang_acc_w). Shape: (N, 1, 3).""" + return self.body_com_ang_acc_w + + # -- CoM in body frame -- + + @property + def com_pos_b(self) -> torch.Tensor: + """CoM position in body frame. Shape: (N, 3).""" + return self.body_com_pose_b[:, 0, :3] + + @property + def com_quat_b(self) -> torch.Tensor: + """CoM orientation in body frame. Shape: (N, 4).""" + return self.body_com_pose_b[:, 0, 3:7] + + # -- Update method -- + + def update(self, dt: float) -> None: + """Update data (no-op for mock).""" + pass + + # -- Setters -- + + def set_default_root_pose(self, value: torch.Tensor) -> None: + self._default_root_pose = value.to(self._device) + + def set_default_root_vel(self, value: torch.Tensor) -> None: + self._default_root_vel = value.to(self._device) + + def set_root_link_pose_w(self, value: torch.Tensor) -> None: + self._root_link_pose_w = value.to(self._device) + + def set_root_link_vel_w(self, value: torch.Tensor) -> None: + self._root_link_vel_w = value.to(self._device) + + def set_root_com_pose_w(self, value: torch.Tensor) -> None: + self._root_com_pose_w = value.to(self._device) + + def set_root_com_vel_w(self, value: torch.Tensor) -> None: + self._root_com_vel_w = value.to(self._device) + + def set_body_link_pose_w(self, value: torch.Tensor) -> None: + self._body_link_pose_w = value.to(self._device) + + def set_body_link_vel_w(self, value: torch.Tensor) -> None: + self._body_link_vel_w = value.to(self._device) + + def set_body_com_pose_w(self, value: torch.Tensor) -> None: + self._body_com_pose_w = value.to(self._device) + + def set_body_com_vel_w(self, value: torch.Tensor) -> None: + self._body_com_vel_w = value.to(self._device) + + def set_body_com_acc_w(self, value: torch.Tensor) -> None: + self._body_com_acc_w = value.to(self._device) + + def set_body_com_pose_b(self, value: torch.Tensor) -> None: + self._body_com_pose_b = value.to(self._device) + + def set_body_mass(self, value: torch.Tensor) -> None: + self._body_mass = value.to(self._device) + + def set_body_inertia(self, value: torch.Tensor) -> None: + self._body_inertia = value.to(self._device) + + def set_mock_data(self, **kwargs) -> None: + """Bulk setter for mock data.""" + for key, value in kwargs.items(): + setter = getattr(self, f"set_{key}", None) + if setter is not None: + setter(value) + else: + raise ValueError(f"Unknown property: {key}") + + +class MockRigidObject: + """Mock rigid object asset for testing without Isaac Sim. + + This class mimics the interface of BaseRigidObject for testing purposes. + It provides the same properties and methods but without simulation dependencies. + """ + + def __init__( + self, + num_instances: int, + body_names: list[str] | None = None, + device: str = "cpu", + ): + """Initialize mock rigid object. + + Args: + num_instances: Number of rigid object instances. + body_names: Names of bodies (single body for rigid object). + device: Device for tensor allocation. + """ + self._num_instances = num_instances + self._num_bodies = 1 + self._device = device + self._body_names = body_names or ["body"] + + self._data = MockRigidObjectData( + num_instances=num_instances, + body_names=self._body_names, + device=device, + ) + + # -- Properties -- + + @property + def data(self) -> MockRigidObjectData: + """Data container for the rigid object.""" + return self._data + + @property + def num_instances(self) -> int: + """Number of rigid object instances.""" + return self._num_instances + + @property + def num_bodies(self) -> int: + """Number of bodies (always 1 for rigid object).""" + return self._num_bodies + + @property + def body_names(self) -> list[str]: + """Body names.""" + return self._body_names + + @property + def device(self) -> str: + """Device for tensor allocation.""" + return self._device + + @property + def root_view(self) -> None: + """Returns None (no physics view in mock).""" + return None + + @property + def instantaneous_wrench_composer(self) -> None: + """Returns None (no wrench composer in mock).""" + return None + + @property + def permanent_wrench_composer(self) -> None: + """Returns None (no wrench composer in mock).""" + return None + + # -- Core methods -- + + def reset(self, env_ids: Sequence[int] | None = None) -> None: + """Reset rigid object state for specified environments.""" + pass + + def write_data_to_sim(self) -> None: + """Write data to simulation (no-op for mock).""" + pass + + def update(self, dt: float) -> None: + """Update rigid object data.""" + self._data.update(dt) + + # -- Finder methods -- + + def find_bodies(self, name_keys: str | Sequence[str], preserve_order: bool = False) -> tuple[list[int], list[str]]: + """Find bodies by name regex patterns.""" + if isinstance(name_keys, str): + name_keys = [name_keys] + + matched_indices = [] + matched_names = [] + + if preserve_order: + for key in name_keys: + pattern = re.compile(key) + for i, name in enumerate(self._body_names): + if pattern.fullmatch(name) and i not in matched_indices: + matched_indices.append(i) + matched_names.append(name) + else: + for i, name in enumerate(self._body_names): + for key in name_keys: + pattern = re.compile(key) + if pattern.fullmatch(name): + matched_indices.append(i) + matched_names.append(name) + break + + return matched_indices, matched_names + + # -- State writer methods (no-op for mock) -- + + def write_root_state_to_sim( + self, + root_state: torch.Tensor, + env_ids: Sequence[int] | None = None, + ) -> None: + """Write root state to simulation.""" + pass + + def write_root_com_state_to_sim( + self, + root_state: torch.Tensor, + env_ids: Sequence[int] | None = None, + ) -> None: + """Write root CoM state to simulation.""" + pass + + def write_root_link_state_to_sim( + self, + root_state: torch.Tensor, + env_ids: Sequence[int] | None = None, + ) -> None: + """Write root link state to simulation.""" + pass + + def write_root_pose_to_sim( + self, + root_pose: torch.Tensor, + env_ids: Sequence[int] | None = None, + ) -> None: + """Write root pose to simulation.""" + pass + + def write_root_link_pose_to_sim( + self, + root_pose: torch.Tensor, + env_ids: Sequence[int] | None = None, + ) -> None: + """Write root link pose to simulation.""" + pass + + def write_root_com_pose_to_sim( + self, + root_pose: torch.Tensor, + env_ids: Sequence[int] | None = None, + ) -> None: + """Write root CoM pose to simulation.""" + pass + + def write_root_velocity_to_sim( + self, + root_velocity: torch.Tensor, + env_ids: Sequence[int] | None = None, + ) -> None: + """Write root velocity to simulation.""" + pass + + def write_root_com_velocity_to_sim( + self, + root_velocity: torch.Tensor, + env_ids: Sequence[int] | None = None, + ) -> None: + """Write root CoM velocity to simulation.""" + pass + + def write_root_link_velocity_to_sim( + self, + root_velocity: torch.Tensor, + env_ids: Sequence[int] | None = None, + ) -> None: + """Write root link velocity to simulation.""" + pass + + # -- Setter methods -- + + def set_masses( + self, + masses: torch.Tensor, + body_ids: Sequence[int] | slice | None = None, + env_ids: Sequence[int] | None = None, + ) -> None: + """Set body masses.""" + pass + + def set_coms( + self, + coms: torch.Tensor, + body_ids: Sequence[int] | slice | None = None, + env_ids: Sequence[int] | None = None, + ) -> None: + """Set body centers of mass.""" + pass + + def set_inertias( + self, + inertias: torch.Tensor, + body_ids: Sequence[int] | slice | None = None, + env_ids: Sequence[int] | None = None, + ) -> None: + """Set body inertias.""" + pass + + def set_external_force_and_torque( + self, + forces: torch.Tensor, + torques: torch.Tensor, + positions: torch.Tensor | None = None, + body_ids: Sequence[int] | slice | None = None, + env_ids: Sequence[int] | None = None, + is_global: bool = True, + ) -> None: + """Set external forces and torques.""" + pass diff --git a/source/isaaclab/isaaclab/test/mock_interfaces/assets/mock_rigid_object_collection.py b/source/isaaclab/isaaclab/test/mock_interfaces/assets/mock_rigid_object_collection.py new file mode 100644 index 00000000000..342bef51c99 --- /dev/null +++ b/source/isaaclab/isaaclab/test/mock_interfaces/assets/mock_rigid_object_collection.py @@ -0,0 +1,579 @@ +# 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 + +"""Mock rigid object collection asset for testing without Isaac Sim.""" + +from __future__ import annotations + +import re +from collections.abc import Sequence + +import torch + + +class MockRigidObjectCollectionData: + """Mock data container for rigid object collection asset. + + This class mimics the interface of BaseRigidObjectCollectionData for testing purposes. + All tensor properties return zero tensors with correct shapes if not explicitly set. + """ + + def __init__( + self, + num_instances: int, + num_bodies: int, + body_names: list[str] | None = None, + device: str = "cpu", + ): + """Initialize mock rigid object collection data. + + Args: + num_instances: Number of environment instances. + num_bodies: Number of bodies in the collection. + body_names: Names of bodies. + device: Device for tensor allocation. + """ + self._num_instances = num_instances + self._num_bodies = num_bodies + self._device = device + + self.body_names = body_names or [f"body_{i}" for i in range(num_bodies)] + + # Default states + self._default_body_pose: torch.Tensor | None = None + self._default_body_vel: torch.Tensor | None = None + + # Body state (link frame) + self._body_link_pose_w: torch.Tensor | None = None + self._body_link_vel_w: torch.Tensor | None = None + + # Body state (CoM frame) + self._body_com_pose_w: torch.Tensor | None = None + self._body_com_vel_w: torch.Tensor | None = None + self._body_com_acc_w: torch.Tensor | None = None + self._body_com_pose_b: torch.Tensor | None = None + + # Body properties + self._body_mass: torch.Tensor | None = None + self._body_inertia: torch.Tensor | None = None + + @property + def device(self) -> str: + """Device for tensor allocation.""" + return self._device + + def _identity_quat(self, *shape: int) -> torch.Tensor: + """Create identity quaternion tensor (w, x, y, z) = (1, 0, 0, 0).""" + quat = torch.zeros(*shape, 4, device=self._device) + quat[..., 0] = 1.0 + return quat + + # -- Default state properties -- + + @property + def default_body_pose(self) -> torch.Tensor: + """Default body poses. Shape: (N, num_bodies, 7).""" + if self._default_body_pose is None: + pose = torch.zeros(self._num_instances, self._num_bodies, 7, device=self._device) + pose[..., 3:7] = self._identity_quat(self._num_instances, self._num_bodies) + return pose + return self._default_body_pose + + @property + def default_body_vel(self) -> torch.Tensor: + """Default body velocities. Shape: (N, num_bodies, 6).""" + if self._default_body_vel is None: + return torch.zeros(self._num_instances, self._num_bodies, 6, device=self._device) + return self._default_body_vel + + @property + def default_body_state(self) -> torch.Tensor: + """Default body states. Shape: (N, num_bodies, 13).""" + return torch.cat([self.default_body_pose, self.default_body_vel], dim=-1) + + # -- Body state properties (link frame) -- + + @property + def body_link_pose_w(self) -> torch.Tensor: + """Body link poses in world frame. Shape: (N, num_bodies, 7).""" + if self._body_link_pose_w is None: + pose = torch.zeros(self._num_instances, self._num_bodies, 7, device=self._device) + pose[..., 3:7] = self._identity_quat(self._num_instances, self._num_bodies) + return pose + return self._body_link_pose_w + + @property + def body_link_vel_w(self) -> torch.Tensor: + """Body link velocities in world frame. Shape: (N, num_bodies, 6).""" + if self._body_link_vel_w is None: + return torch.zeros(self._num_instances, self._num_bodies, 6, device=self._device) + return self._body_link_vel_w + + @property + def body_link_state_w(self) -> torch.Tensor: + """Body link states in world frame. Shape: (N, num_bodies, 13).""" + return torch.cat([self.body_link_pose_w, self.body_link_vel_w], dim=-1) + + @property + def body_link_pos_w(self) -> torch.Tensor: + """Body link positions. Shape: (N, num_bodies, 3).""" + return self.body_link_pose_w[..., :3] + + @property + def body_link_quat_w(self) -> torch.Tensor: + """Body link orientations. Shape: (N, num_bodies, 4).""" + return self.body_link_pose_w[..., 3:7] + + @property + def body_link_lin_vel_w(self) -> torch.Tensor: + """Body link linear velocities. Shape: (N, num_bodies, 3).""" + return self.body_link_vel_w[..., :3] + + @property + def body_link_ang_vel_w(self) -> torch.Tensor: + """Body link angular velocities. Shape: (N, num_bodies, 3).""" + return self.body_link_vel_w[..., 3:6] + + # -- Body state properties (CoM frame) -- + + @property + def body_com_pose_w(self) -> torch.Tensor: + """Body CoM poses in world frame. Shape: (N, num_bodies, 7).""" + if self._body_com_pose_w is None: + return self.body_link_pose_w.clone() + return self._body_com_pose_w + + @property + def body_com_vel_w(self) -> torch.Tensor: + """Body CoM velocities in world frame. Shape: (N, num_bodies, 6).""" + if self._body_com_vel_w is None: + return self.body_link_vel_w.clone() + return self._body_com_vel_w + + @property + def body_com_state_w(self) -> torch.Tensor: + """Body CoM states in world frame. Shape: (N, num_bodies, 13).""" + return torch.cat([self.body_com_pose_w, self.body_com_vel_w], dim=-1) + + @property + def body_com_acc_w(self) -> torch.Tensor: + """Body CoM accelerations in world frame. Shape: (N, num_bodies, 6).""" + if self._body_com_acc_w is None: + return torch.zeros(self._num_instances, self._num_bodies, 6, device=self._device) + return self._body_com_acc_w + + @property + def body_com_pose_b(self) -> torch.Tensor: + """Body CoM poses in body frame. Shape: (N, num_bodies, 7).""" + if self._body_com_pose_b is None: + pose = torch.zeros(self._num_instances, self._num_bodies, 7, device=self._device) + pose[..., 3:7] = self._identity_quat(self._num_instances, self._num_bodies) + return pose + return self._body_com_pose_b + + @property + def body_com_pos_w(self) -> torch.Tensor: + """Body CoM positions. Shape: (N, num_bodies, 3).""" + return self.body_com_pose_w[..., :3] + + @property + def body_com_quat_w(self) -> torch.Tensor: + """Body CoM orientations. Shape: (N, num_bodies, 4).""" + return self.body_com_pose_w[..., 3:7] + + @property + def body_com_lin_vel_w(self) -> torch.Tensor: + """Body CoM linear velocities. Shape: (N, num_bodies, 3).""" + return self.body_com_vel_w[..., :3] + + @property + def body_com_ang_vel_w(self) -> torch.Tensor: + """Body CoM angular velocities. Shape: (N, num_bodies, 3).""" + return self.body_com_vel_w[..., 3:6] + + @property + def body_com_lin_acc_w(self) -> torch.Tensor: + """Body CoM linear accelerations. Shape: (N, num_bodies, 3).""" + return self.body_com_acc_w[..., :3] + + @property + def body_com_ang_acc_w(self) -> torch.Tensor: + """Body CoM angular accelerations. Shape: (N, num_bodies, 3).""" + return self.body_com_acc_w[..., 3:6] + + @property + def body_com_pos_b(self) -> torch.Tensor: + """Body CoM positions in body frame. Shape: (N, num_bodies, 3).""" + return self.body_com_pose_b[..., :3] + + @property + def body_com_quat_b(self) -> torch.Tensor: + """Body CoM orientations in body frame. Shape: (N, num_bodies, 4).""" + return self.body_com_pose_b[..., 3:7] + + # -- Body properties -- + + @property + def body_mass(self) -> torch.Tensor: + """Body masses. Shape: (N, num_bodies).""" + if self._body_mass is None: + return torch.ones(self._num_instances, self._num_bodies, device=self._device) + return self._body_mass + + @property + def body_inertia(self) -> torch.Tensor: + """Body inertias (flattened 3x3). Shape: (N, num_bodies, 9).""" + if self._body_inertia is None: + inertia = torch.zeros(self._num_instances, self._num_bodies, 9, device=self._device) + inertia[..., 0] = 1.0 # Ixx + inertia[..., 4] = 1.0 # Iyy + inertia[..., 8] = 1.0 # Izz + return inertia + return self._body_inertia + + # -- Derived properties -- + + @property + def projected_gravity_b(self) -> torch.Tensor: + """Gravity projection on bodies. Shape: (N, num_bodies, 3).""" + gravity = torch.zeros(self._num_instances, self._num_bodies, 3, device=self._device) + gravity[..., 2] = -1.0 + return gravity + + @property + def heading_w(self) -> torch.Tensor: + """Yaw heading per body. Shape: (N, num_bodies).""" + return torch.zeros(self._num_instances, self._num_bodies, device=self._device) + + @property + def body_link_lin_vel_b(self) -> torch.Tensor: + """Body link linear velocities in body frame. Shape: (N, num_bodies, 3).""" + return self.body_link_lin_vel_w.clone() + + @property + def body_link_ang_vel_b(self) -> torch.Tensor: + """Body link angular velocities in body frame. Shape: (N, num_bodies, 3).""" + return self.body_link_ang_vel_w.clone() + + @property + def body_com_lin_vel_b(self) -> torch.Tensor: + """Body CoM linear velocities in body frame. Shape: (N, num_bodies, 3).""" + return self.body_com_lin_vel_w.clone() + + @property + def body_com_ang_vel_b(self) -> torch.Tensor: + """Body CoM angular velocities in body frame. Shape: (N, num_bodies, 3).""" + return self.body_com_ang_vel_w.clone() + + # -- Convenience alias -- + + @property + def body_state_w(self) -> torch.Tensor: + """Body states (alias for body_com_state_w). Shape: (N, num_bodies, 13).""" + return self.body_com_state_w + + # -- Update method -- + + def update(self, dt: float) -> None: + """Update data (no-op for mock).""" + pass + + # -- Setters -- + + def set_default_body_pose(self, value: torch.Tensor) -> None: + self._default_body_pose = value.to(self._device) + + def set_default_body_vel(self, value: torch.Tensor) -> None: + self._default_body_vel = value.to(self._device) + + def set_body_link_pose_w(self, value: torch.Tensor) -> None: + self._body_link_pose_w = value.to(self._device) + + def set_body_link_vel_w(self, value: torch.Tensor) -> None: + self._body_link_vel_w = value.to(self._device) + + def set_body_com_pose_w(self, value: torch.Tensor) -> None: + self._body_com_pose_w = value.to(self._device) + + def set_body_com_vel_w(self, value: torch.Tensor) -> None: + self._body_com_vel_w = value.to(self._device) + + def set_body_com_acc_w(self, value: torch.Tensor) -> None: + self._body_com_acc_w = value.to(self._device) + + def set_body_com_pose_b(self, value: torch.Tensor) -> None: + self._body_com_pose_b = value.to(self._device) + + def set_body_mass(self, value: torch.Tensor) -> None: + self._body_mass = value.to(self._device) + + def set_body_inertia(self, value: torch.Tensor) -> None: + self._body_inertia = value.to(self._device) + + def set_mock_data(self, **kwargs) -> None: + """Bulk setter for mock data.""" + for key, value in kwargs.items(): + setter = getattr(self, f"set_{key}", None) + if setter is not None: + setter(value) + else: + raise ValueError(f"Unknown property: {key}") + + +class MockRigidObjectCollection: + """Mock rigid object collection asset for testing without Isaac Sim. + + This class mimics the interface of BaseRigidObjectCollection for testing purposes. + It provides the same properties and methods but without simulation dependencies. + """ + + def __init__( + self, + num_instances: int, + num_bodies: int, + body_names: list[str] | None = None, + device: str = "cpu", + ): + """Initialize mock rigid object collection. + + Args: + num_instances: Number of environment instances. + num_bodies: Number of bodies in the collection. + body_names: Names of bodies. + device: Device for tensor allocation. + """ + self._num_instances = num_instances + self._num_bodies = num_bodies + self._device = device + self._body_names = body_names or [f"body_{i}" for i in range(num_bodies)] + + self._data = MockRigidObjectCollectionData( + num_instances=num_instances, + num_bodies=num_bodies, + body_names=self._body_names, + device=device, + ) + + # -- Properties -- + + @property + def data(self) -> MockRigidObjectCollectionData: + """Data container for the rigid object collection.""" + return self._data + + @property + def num_instances(self) -> int: + """Number of environment instances.""" + return self._num_instances + + @property + def num_bodies(self) -> int: + """Number of bodies in the collection.""" + return self._num_bodies + + @property + def body_names(self) -> list[str]: + """Body names.""" + return self._body_names + + @property + def device(self) -> str: + """Device for tensor allocation.""" + return self._device + + @property + def root_view(self) -> None: + """Returns None (no physics view in mock).""" + return None + + @property + def instantaneous_wrench_composer(self) -> None: + """Returns None (no wrench composer in mock).""" + return None + + @property + def permanent_wrench_composer(self) -> None: + """Returns None (no wrench composer in mock).""" + return None + + # -- Core methods -- + + def reset( + self, + env_ids: Sequence[int] | None = None, + object_ids: slice | torch.Tensor | None = None, + ) -> None: + """Reset rigid object collection state for specified environments.""" + pass + + def write_data_to_sim(self) -> None: + """Write data to simulation (no-op for mock).""" + pass + + def update(self, dt: float) -> None: + """Update rigid object collection data.""" + self._data.update(dt) + + # -- Finder methods -- + + def find_bodies( + self, name_keys: str | Sequence[str], preserve_order: bool = False + ) -> tuple[torch.Tensor, list[str], list[int]]: + """Find bodies by name regex patterns. + + Returns: + Tuple of (body_mask, body_names, body_indices). + """ + if isinstance(name_keys, str): + name_keys = [name_keys] + + matched_indices = [] + matched_names = [] + + if preserve_order: + for key in name_keys: + pattern = re.compile(key) + for i, name in enumerate(self._body_names): + if pattern.fullmatch(name) and i not in matched_indices: + matched_indices.append(i) + matched_names.append(name) + else: + for i, name in enumerate(self._body_names): + for key in name_keys: + pattern = re.compile(key) + if pattern.fullmatch(name): + matched_indices.append(i) + matched_names.append(name) + break + + # Create body mask + body_mask = torch.zeros(self._num_bodies, dtype=torch.bool, device=self._device) + body_mask[matched_indices] = True + + return body_mask, matched_names, matched_indices + + # -- State writer methods (no-op for mock) -- + + def write_body_state_to_sim( + self, + body_states: torch.Tensor, + env_ids: Sequence[int] | None = None, + body_ids: Sequence[int] | slice | None = None, + ) -> None: + """Write body states to simulation.""" + pass + + def write_body_com_state_to_sim( + self, + body_states: torch.Tensor, + env_ids: Sequence[int] | None = None, + body_ids: Sequence[int] | slice | None = None, + ) -> None: + """Write body CoM states to simulation.""" + pass + + def write_body_link_state_to_sim( + self, + body_states: torch.Tensor, + env_ids: Sequence[int] | None = None, + body_ids: Sequence[int] | slice | None = None, + ) -> None: + """Write body link states to simulation.""" + pass + + def write_body_pose_to_sim( + self, + body_poses: torch.Tensor, + env_ids: Sequence[int] | None = None, + body_ids: Sequence[int] | slice | None = None, + ) -> None: + """Write body poses to simulation.""" + pass + + def write_body_link_pose_to_sim( + self, + body_poses: torch.Tensor, + env_ids: Sequence[int] | None = None, + body_ids: Sequence[int] | slice | None = None, + ) -> None: + """Write body link poses to simulation.""" + pass + + def write_body_com_pose_to_sim( + self, + body_poses: torch.Tensor, + env_ids: Sequence[int] | None = None, + body_ids: Sequence[int] | slice | None = None, + ) -> None: + """Write body CoM poses to simulation.""" + pass + + def write_body_velocity_to_sim( + self, + body_velocities: torch.Tensor, + env_ids: Sequence[int] | None = None, + body_ids: Sequence[int] | slice | None = None, + ) -> None: + """Write body velocities to simulation.""" + pass + + def write_body_com_velocity_to_sim( + self, + body_velocities: torch.Tensor, + env_ids: Sequence[int] | None = None, + body_ids: Sequence[int] | slice | None = None, + ) -> None: + """Write body CoM velocities to simulation.""" + pass + + def write_body_link_velocity_to_sim( + self, + body_velocities: torch.Tensor, + env_ids: Sequence[int] | None = None, + body_ids: Sequence[int] | slice | None = None, + ) -> None: + """Write body link velocities to simulation.""" + pass + + # -- Setter methods -- + + def set_masses( + self, + masses: torch.Tensor, + body_ids: Sequence[int] | slice | None = None, + env_ids: Sequence[int] | None = None, + ) -> None: + """Set body masses.""" + pass + + def set_coms( + self, + coms: torch.Tensor, + body_ids: Sequence[int] | slice | None = None, + env_ids: Sequence[int] | None = None, + ) -> None: + """Set body centers of mass.""" + pass + + def set_inertias( + self, + inertias: torch.Tensor, + body_ids: Sequence[int] | slice | None = None, + env_ids: Sequence[int] | None = None, + ) -> None: + """Set body inertias.""" + pass + + def set_external_force_and_torque( + self, + forces: torch.Tensor, + torques: torch.Tensor, + positions: torch.Tensor | None = None, + body_ids: Sequence[int] | slice | None = None, + env_ids: Sequence[int] | None = None, + is_global: bool = True, + ) -> None: + """Set external forces and torques.""" + pass diff --git a/source/isaaclab/isaaclab/test/mock_interfaces/sensors/__init__.py b/source/isaaclab/isaaclab/test/mock_interfaces/sensors/__init__.py new file mode 100644 index 00000000000..8ea858817fc --- /dev/null +++ b/source/isaaclab/isaaclab/test/mock_interfaces/sensors/__init__.py @@ -0,0 +1,16 @@ +# 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 + +"""Mock sensor interfaces for testing without Isaac Sim.""" + +from .mock_contact_sensor import MockContactSensor, MockContactSensorData +from .mock_frame_transformer import MockFrameTransformer, MockFrameTransformerData +from .mock_imu import MockImu, MockImuData +from .factories import ( + create_mock_contact_sensor, + create_mock_foot_contact_sensor, + create_mock_frame_transformer, + create_mock_imu, +) diff --git a/source/isaaclab/isaaclab/test/mock_interfaces/sensors/factories.py b/source/isaaclab/isaaclab/test/mock_interfaces/sensors/factories.py new file mode 100644 index 00000000000..8f99a4854ad --- /dev/null +++ b/source/isaaclab/isaaclab/test/mock_interfaces/sensors/factories.py @@ -0,0 +1,121 @@ +# 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 + +"""Factory functions for creating pre-configured mock sensors.""" + +from __future__ import annotations + +import torch + +from .mock_contact_sensor import MockContactSensor +from .mock_frame_transformer import MockFrameTransformer +from .mock_imu import MockImu + + +def create_mock_imu( + num_instances: int = 1, + device: str = "cpu", + gravity: tuple[float, float, float] = (0.0, 0.0, -1.0), +) -> MockImu: + """Create a mock IMU sensor with default configuration. + + Args: + num_instances: Number of sensor instances. + device: Device for tensor allocation. + gravity: Default gravity direction in body frame. + + Returns: + Configured MockImu instance. + """ + imu = MockImu(num_instances=num_instances, device=device) + imu.data.set_projected_gravity_b(torch.tensor([gravity], device=device).expand(num_instances, -1).clone()) + return imu + + +def create_mock_contact_sensor( + num_instances: int = 1, + num_bodies: int = 1, + body_names: list[str] | None = None, + device: str = "cpu", + history_length: int = 0, + num_filter_bodies: int = 0, +) -> MockContactSensor: + """Create a mock contact sensor with default configuration. + + Args: + num_instances: Number of environment instances. + num_bodies: Number of bodies with contact sensors. + body_names: Names of bodies with contact sensors. + device: Device for tensor allocation. + history_length: Length of history buffer for forces. + num_filter_bodies: Number of filter bodies for force matrix. + + Returns: + Configured MockContactSensor instance. + """ + return MockContactSensor( + num_instances=num_instances, + num_bodies=num_bodies, + body_names=body_names, + device=device, + history_length=history_length, + num_filter_bodies=num_filter_bodies, + ) + + +def create_mock_foot_contact_sensor( + num_instances: int = 1, + num_feet: int = 4, + foot_names: list[str] | None = None, + device: str = "cpu", + history_length: int = 0, +) -> MockContactSensor: + """Create a mock foot contact sensor for quadruped robots. + + Args: + num_instances: Number of environment instances. + num_feet: Number of feet (default 4 for quadruped). + foot_names: Names of feet. Defaults to FL, FR, RL, RR. + device: Device for tensor allocation. + history_length: Length of history buffer for forces. + + Returns: + Configured MockContactSensor instance for foot contacts. + """ + if foot_names is None: + foot_names = ["FL_foot", "FR_foot", "RL_foot", "RR_foot"][:num_feet] + + return MockContactSensor( + num_instances=num_instances, + num_bodies=num_feet, + body_names=foot_names, + device=device, + history_length=history_length, + ) + + +def create_mock_frame_transformer( + num_instances: int = 1, + num_target_frames: int = 1, + target_frame_names: list[str] | None = None, + device: str = "cpu", +) -> MockFrameTransformer: + """Create a mock frame transformer sensor with default configuration. + + Args: + num_instances: Number of environment instances. + num_target_frames: Number of target frames to track. + target_frame_names: Names of target frames. + device: Device for tensor allocation. + + Returns: + Configured MockFrameTransformer instance. + """ + return MockFrameTransformer( + num_instances=num_instances, + num_target_frames=num_target_frames, + target_frame_names=target_frame_names, + device=device, + ) diff --git a/source/isaaclab/isaaclab/test/mock_interfaces/sensors/mock_contact_sensor.py b/source/isaaclab/isaaclab/test/mock_interfaces/sensors/mock_contact_sensor.py new file mode 100644 index 00000000000..140d5cc5d8b --- /dev/null +++ b/source/isaaclab/isaaclab/test/mock_interfaces/sensors/mock_contact_sensor.py @@ -0,0 +1,419 @@ +# 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 + +"""Mock contact sensor for testing without Isaac Sim.""" + +from __future__ import annotations + +import re +from collections.abc import Sequence + +import torch + + +class MockContactSensorData: + """Mock data container for contact sensor. + + This class mimics the interface of BaseContactSensorData for testing purposes. + All tensor properties return zero tensors with correct shapes if not explicitly set. + """ + + def __init__( + self, + num_instances: int, + num_bodies: int, + device: str = "cpu", + history_length: int = 0, + num_filter_bodies: int = 0, + ): + """Initialize mock contact sensor data. + + Args: + num_instances: Number of environment instances. + num_bodies: Number of bodies with contact sensors. + device: Device for tensor allocation. + history_length: Length of history buffer for forces. + num_filter_bodies: Number of filter bodies for force matrix. + """ + self._num_instances = num_instances + self._num_bodies = num_bodies + self._device = device + self._history_length = history_length + self._num_filter_bodies = num_filter_bodies + + # Internal storage for mock data + self._pos_w: torch.Tensor | None = None + self._quat_w: torch.Tensor | None = None + self._net_forces_w: torch.Tensor | None = None + self._net_forces_w_history: torch.Tensor | None = None + self._force_matrix_w: torch.Tensor | None = None + self._force_matrix_w_history: torch.Tensor | None = None + self._contact_pos_w: torch.Tensor | None = None + self._friction_forces_w: torch.Tensor | None = None + self._last_air_time: torch.Tensor | None = None + self._current_air_time: torch.Tensor | None = None + self._last_contact_time: torch.Tensor | None = None + self._current_contact_time: torch.Tensor | None = None + + # -- Properties -- + + @property + def pos_w(self) -> torch.Tensor | None: + """Position of sensor origins in world frame. Shape: (N, B, 3).""" + if self._pos_w is None: + return torch.zeros(self._num_instances, self._num_bodies, 3, device=self._device) + return self._pos_w + + @property + def quat_w(self) -> torch.Tensor | None: + """Orientation (w, x, y, z) in world frame. Shape: (N, B, 4).""" + if self._quat_w is None: + # Default to identity quaternion + quat = torch.zeros(self._num_instances, self._num_bodies, 4, device=self._device) + quat[..., 0] = 1.0 + return quat + return self._quat_w + + @property + def pose_w(self) -> torch.Tensor | None: + """Pose in world frame (pos + quat). Shape: (N, B, 7).""" + return torch.cat([self.pos_w, self.quat_w], dim=-1) + + @property + def net_forces_w(self) -> torch.Tensor: + """Net normal contact forces in world frame. Shape: (N, B, 3).""" + if self._net_forces_w is None: + return torch.zeros(self._num_instances, self._num_bodies, 3, device=self._device) + return self._net_forces_w + + @property + def net_forces_w_history(self) -> torch.Tensor | None: + """History of net forces. Shape: (N, T, B, 3).""" + if self._history_length == 0: + return None + if self._net_forces_w_history is None: + return torch.zeros(self._num_instances, self._history_length, self._num_bodies, 3, device=self._device) + return self._net_forces_w_history + + @property + def force_matrix_w(self) -> torch.Tensor | None: + """Filtered contact forces. Shape: (N, B, M, 3).""" + if self._num_filter_bodies == 0: + return None + if self._force_matrix_w is None: + return torch.zeros(self._num_instances, self._num_bodies, self._num_filter_bodies, 3, device=self._device) + return self._force_matrix_w + + @property + def force_matrix_w_history(self) -> torch.Tensor | None: + """History of filtered forces. Shape: (N, T, B, M, 3).""" + if self._history_length == 0 or self._num_filter_bodies == 0: + return None + if self._force_matrix_w_history is None: + return torch.zeros( + self._num_instances, + self._history_length, + self._num_bodies, + self._num_filter_bodies, + 3, + device=self._device, + ) + return self._force_matrix_w_history + + @property + def contact_pos_w(self) -> torch.Tensor | None: + """Contact point positions in world frame. Shape: (N, B, M, 3).""" + if self._num_filter_bodies == 0: + return None + if self._contact_pos_w is None: + return torch.zeros(self._num_instances, self._num_bodies, self._num_filter_bodies, 3, device=self._device) + return self._contact_pos_w + + @property + def friction_forces_w(self) -> torch.Tensor | None: + """Friction forces in world frame. Shape: (N, B, M, 3).""" + if self._num_filter_bodies == 0: + return None + if self._friction_forces_w is None: + return torch.zeros(self._num_instances, self._num_bodies, self._num_filter_bodies, 3, device=self._device) + return self._friction_forces_w + + @property + def last_air_time(self) -> torch.Tensor: + """Time in air before last contact. Shape: (N, B).""" + if self._last_air_time is None: + return torch.zeros(self._num_instances, self._num_bodies, device=self._device) + return self._last_air_time + + @property + def current_air_time(self) -> torch.Tensor: + """Current time in air. Shape: (N, B).""" + if self._current_air_time is None: + return torch.zeros(self._num_instances, self._num_bodies, device=self._device) + return self._current_air_time + + @property + def last_contact_time(self) -> torch.Tensor: + """Time in contact before last detach. Shape: (N, B).""" + if self._last_contact_time is None: + return torch.zeros(self._num_instances, self._num_bodies, device=self._device) + return self._last_contact_time + + @property + def current_contact_time(self) -> torch.Tensor: + """Current time in contact. Shape: (N, B).""" + if self._current_contact_time is None: + return torch.zeros(self._num_instances, self._num_bodies, device=self._device) + return self._current_contact_time + + # -- Setters -- + + def set_pos_w(self, value: torch.Tensor) -> None: + """Set position in world frame.""" + self._pos_w = value.to(self._device) + + def set_quat_w(self, value: torch.Tensor) -> None: + """Set orientation in world frame.""" + self._quat_w = value.to(self._device) + + def set_net_forces_w(self, value: torch.Tensor) -> None: + """Set net contact forces.""" + self._net_forces_w = value.to(self._device) + + def set_net_forces_w_history(self, value: torch.Tensor) -> None: + """Set net forces history.""" + self._net_forces_w_history = value.to(self._device) + + def set_force_matrix_w(self, value: torch.Tensor) -> None: + """Set filtered contact forces.""" + self._force_matrix_w = value.to(self._device) + + def set_force_matrix_w_history(self, value: torch.Tensor) -> None: + """Set filtered forces history.""" + self._force_matrix_w_history = value.to(self._device) + + def set_contact_pos_w(self, value: torch.Tensor) -> None: + """Set contact point positions.""" + self._contact_pos_w = value.to(self._device) + + def set_friction_forces_w(self, value: torch.Tensor) -> None: + """Set friction forces.""" + self._friction_forces_w = value.to(self._device) + + def set_last_air_time(self, value: torch.Tensor) -> None: + """Set last air time.""" + self._last_air_time = value.to(self._device) + + def set_current_air_time(self, value: torch.Tensor) -> None: + """Set current air time.""" + self._current_air_time = value.to(self._device) + + def set_last_contact_time(self, value: torch.Tensor) -> None: + """Set last contact time.""" + self._last_contact_time = value.to(self._device) + + def set_current_contact_time(self, value: torch.Tensor) -> None: + """Set current contact time.""" + self._current_contact_time = value.to(self._device) + + def set_mock_data( + self, + pos_w: torch.Tensor | None = None, + quat_w: torch.Tensor | None = None, + net_forces_w: torch.Tensor | None = None, + net_forces_w_history: torch.Tensor | None = None, + force_matrix_w: torch.Tensor | None = None, + force_matrix_w_history: torch.Tensor | None = None, + contact_pos_w: torch.Tensor | None = None, + friction_forces_w: torch.Tensor | None = None, + last_air_time: torch.Tensor | None = None, + current_air_time: torch.Tensor | None = None, + last_contact_time: torch.Tensor | None = None, + current_contact_time: torch.Tensor | None = None, + ) -> None: + """Bulk setter for mock data. + + Args: + pos_w: Position in world frame. Shape: (N, B, 3). + quat_w: Orientation in world frame. Shape: (N, B, 4). + net_forces_w: Net contact forces. Shape: (N, B, 3). + net_forces_w_history: History of net forces. Shape: (N, T, B, 3). + force_matrix_w: Filtered contact forces. Shape: (N, B, M, 3). + force_matrix_w_history: History of filtered forces. Shape: (N, T, B, M, 3). + contact_pos_w: Contact point positions. Shape: (N, B, M, 3). + friction_forces_w: Friction forces. Shape: (N, B, M, 3). + last_air_time: Time in air before last contact. Shape: (N, B). + current_air_time: Current time in air. Shape: (N, B). + last_contact_time: Time in contact before last detach. Shape: (N, B). + current_contact_time: Current time in contact. Shape: (N, B). + """ + if pos_w is not None: + self.set_pos_w(pos_w) + if quat_w is not None: + self.set_quat_w(quat_w) + if net_forces_w is not None: + self.set_net_forces_w(net_forces_w) + if net_forces_w_history is not None: + self.set_net_forces_w_history(net_forces_w_history) + if force_matrix_w is not None: + self.set_force_matrix_w(force_matrix_w) + if force_matrix_w_history is not None: + self.set_force_matrix_w_history(force_matrix_w_history) + if contact_pos_w is not None: + self.set_contact_pos_w(contact_pos_w) + if friction_forces_w is not None: + self.set_friction_forces_w(friction_forces_w) + if last_air_time is not None: + self.set_last_air_time(last_air_time) + if current_air_time is not None: + self.set_current_air_time(current_air_time) + if last_contact_time is not None: + self.set_last_contact_time(last_contact_time) + if current_contact_time is not None: + self.set_current_contact_time(current_contact_time) + + +class MockContactSensor: + """Mock contact sensor for testing without Isaac Sim. + + This class mimics the interface of BaseContactSensor for testing purposes. + It provides the same properties and methods but without simulation dependencies. + """ + + def __init__( + self, + num_instances: int, + num_bodies: int, + body_names: list[str] | None = None, + device: str = "cpu", + history_length: int = 0, + num_filter_bodies: int = 0, + ): + """Initialize mock contact sensor. + + Args: + num_instances: Number of environment instances. + num_bodies: Number of bodies with contact sensors. + body_names: Names of bodies with contact sensors. + device: Device for tensor allocation. + history_length: Length of history buffer for forces. + num_filter_bodies: Number of filter bodies for force matrix. + """ + self._num_instances = num_instances + self._num_bodies = num_bodies + self._body_names = body_names or [f"body_{i}" for i in range(num_bodies)] + self._device = device + self._data = MockContactSensorData(num_instances, num_bodies, device, history_length, num_filter_bodies) + + # -- Properties -- + + @property + def data(self) -> MockContactSensorData: + """Data container for the sensor.""" + return self._data + + @property + def num_instances(self) -> int: + """Number of sensor instances.""" + return self._num_instances + + @property + def num_bodies(self) -> int: + """Number of bodies with contact sensors.""" + return self._num_bodies + + @property + def body_names(self) -> list[str]: + """Ordered names of bodies with contact sensors.""" + return self._body_names + + @property + def contact_view(self) -> None: + """Returns None (no PhysX view in mock).""" + return None + + @property + def device(self) -> str: + """Device for tensor allocation.""" + return self._device + + # -- Methods -- + + def find_bodies(self, name_keys: str | Sequence[str], preserve_order: bool = False) -> tuple[list[int], list[str]]: + """Find bodies by name regex patterns. + + Args: + name_keys: Regex pattern(s) to match body names. + preserve_order: If True, preserve order of name_keys in output. + + Returns: + Tuple of (body_indices, body_names) matching the patterns. + """ + if isinstance(name_keys, str): + name_keys = [name_keys] + + matched_indices = [] + matched_names = [] + + if preserve_order: + for key in name_keys: + pattern = re.compile(key) + for i, name in enumerate(self._body_names): + if pattern.fullmatch(name) and i not in matched_indices: + matched_indices.append(i) + matched_names.append(name) + else: + for i, name in enumerate(self._body_names): + for key in name_keys: + pattern = re.compile(key) + if pattern.fullmatch(name): + matched_indices.append(i) + matched_names.append(name) + break + + return matched_indices, matched_names + + def compute_first_contact(self, dt: float, abs_tol: float = 1.0e-8) -> torch.Tensor: + """Check which bodies established contact within dt seconds. + + Args: + dt: Time window to check for first contact. + abs_tol: Absolute tolerance for contact time comparison. + + Returns: + Boolean tensor of shape (N, B) indicating first contact. + """ + return self._data.current_contact_time < (dt + abs_tol) + + def compute_first_air(self, dt: float, abs_tol: float = 1.0e-8) -> torch.Tensor: + """Check which bodies broke contact within dt seconds. + + Args: + dt: Time window to check for first air. + abs_tol: Absolute tolerance for air time comparison. + + Returns: + Boolean tensor of shape (N, B) indicating first air. + """ + return self._data.current_air_time < (dt + abs_tol) + + def reset(self, env_ids: Sequence[int] | None = None) -> None: + """Reset sensor state for specified environments. + + Args: + env_ids: Environment indices to reset. If None, resets all. + """ + # No-op for mock - data persists until explicitly changed + pass + + def update(self, dt: float, force_recompute: bool = False) -> None: + """Update sensor. + + Args: + dt: Time step since last update. + force_recompute: Force recomputation of buffers. + """ + # No-op for mock - data is set explicitly + pass diff --git a/source/isaaclab/isaaclab/test/mock_interfaces/sensors/mock_frame_transformer.py b/source/isaaclab/isaaclab/test/mock_interfaces/sensors/mock_frame_transformer.py new file mode 100644 index 00000000000..55fce15edc9 --- /dev/null +++ b/source/isaaclab/isaaclab/test/mock_interfaces/sensors/mock_frame_transformer.py @@ -0,0 +1,288 @@ +# 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 + +"""Mock frame transformer sensor for testing without Isaac Sim.""" + +from __future__ import annotations + +import re +from collections.abc import Sequence + +import torch + + +class MockFrameTransformerData: + """Mock data container for frame transformer sensor. + + This class mimics the interface of BaseFrameTransformerData for testing purposes. + All tensor properties return zero tensors with correct shapes if not explicitly set. + """ + + def __init__( + self, + num_instances: int, + num_target_frames: int, + target_frame_names: list[str] | None = None, + device: str = "cpu", + ): + """Initialize mock frame transformer data. + + Args: + num_instances: Number of environment instances. + num_target_frames: Number of target frames to track. + target_frame_names: Names of target frames. + device: Device for tensor allocation. + """ + self._num_instances = num_instances + self._num_target_frames = num_target_frames + self._target_frame_names = target_frame_names or [f"frame_{i}" for i in range(num_target_frames)] + self._device = device + + # Internal storage for mock data + self._source_pos_w: torch.Tensor | None = None + self._source_quat_w: torch.Tensor | None = None + self._target_pos_w: torch.Tensor | None = None + self._target_quat_w: torch.Tensor | None = None + self._target_pos_source: torch.Tensor | None = None + self._target_quat_source: torch.Tensor | None = None + + # -- Properties -- + + @property + def target_frame_names(self) -> list[str]: + """Names of target frames.""" + return self._target_frame_names + + @property + def source_pos_w(self) -> torch.Tensor: + """Position of source frame in world frame. Shape: (N, 3).""" + if self._source_pos_w is None: + return torch.zeros(self._num_instances, 3, device=self._device) + return self._source_pos_w + + @property + def source_quat_w(self) -> torch.Tensor: + """Orientation (w, x, y, z) of source frame in world frame. Shape: (N, 4).""" + if self._source_quat_w is None: + quat = torch.zeros(self._num_instances, 4, device=self._device) + quat[:, 0] = 1.0 + return quat + return self._source_quat_w + + @property + def source_pose_w(self) -> torch.Tensor: + """Pose of source frame in world frame. Shape: (N, 7).""" + return torch.cat([self.source_pos_w, self.source_quat_w], dim=-1) + + @property + def target_pos_w(self) -> torch.Tensor: + """Position of target frames in world frame. Shape: (N, M, 3).""" + if self._target_pos_w is None: + return torch.zeros(self._num_instances, self._num_target_frames, 3, device=self._device) + return self._target_pos_w + + @property + def target_quat_w(self) -> torch.Tensor: + """Orientation (w, x, y, z) of target frames in world frame. Shape: (N, M, 4).""" + if self._target_quat_w is None: + quat = torch.zeros(self._num_instances, self._num_target_frames, 4, device=self._device) + quat[..., 0] = 1.0 + return quat + return self._target_quat_w + + @property + def target_pose_w(self) -> torch.Tensor: + """Pose of target frames in world frame. Shape: (N, M, 7).""" + return torch.cat([self.target_pos_w, self.target_quat_w], dim=-1) + + @property + def target_pos_source(self) -> torch.Tensor: + """Position of target frames relative to source frame. Shape: (N, M, 3).""" + if self._target_pos_source is None: + return torch.zeros(self._num_instances, self._num_target_frames, 3, device=self._device) + return self._target_pos_source + + @property + def target_quat_source(self) -> torch.Tensor: + """Orientation (w, x, y, z) of target frames relative to source frame. Shape: (N, M, 4).""" + if self._target_quat_source is None: + quat = torch.zeros(self._num_instances, self._num_target_frames, 4, device=self._device) + quat[..., 0] = 1.0 + return quat + return self._target_quat_source + + @property + def target_pose_source(self) -> torch.Tensor: + """Pose of target frames relative to source frame. Shape: (N, M, 7).""" + return torch.cat([self.target_pos_source, self.target_quat_source], dim=-1) + + # -- Setters -- + + def set_source_pos_w(self, value: torch.Tensor) -> None: + """Set source position in world frame.""" + self._source_pos_w = value.to(self._device) + + def set_source_quat_w(self, value: torch.Tensor) -> None: + """Set source orientation in world frame.""" + self._source_quat_w = value.to(self._device) + + def set_target_pos_w(self, value: torch.Tensor) -> None: + """Set target positions in world frame.""" + self._target_pos_w = value.to(self._device) + + def set_target_quat_w(self, value: torch.Tensor) -> None: + """Set target orientations in world frame.""" + self._target_quat_w = value.to(self._device) + + def set_target_pos_source(self, value: torch.Tensor) -> None: + """Set target positions relative to source.""" + self._target_pos_source = value.to(self._device) + + def set_target_quat_source(self, value: torch.Tensor) -> None: + """Set target orientations relative to source.""" + self._target_quat_source = value.to(self._device) + + def set_mock_data( + self, + source_pos_w: torch.Tensor | None = None, + source_quat_w: torch.Tensor | None = None, + target_pos_w: torch.Tensor | None = None, + target_quat_w: torch.Tensor | None = None, + target_pos_source: torch.Tensor | None = None, + target_quat_source: torch.Tensor | None = None, + ) -> None: + """Bulk setter for mock data. + + Args: + source_pos_w: Source position in world frame. Shape: (N, 3). + source_quat_w: Source orientation in world frame. Shape: (N, 4). + target_pos_w: Target positions in world frame. Shape: (N, M, 3). + target_quat_w: Target orientations in world frame. Shape: (N, M, 4). + target_pos_source: Target positions relative to source. Shape: (N, M, 3). + target_quat_source: Target orientations relative to source. Shape: (N, M, 4). + """ + if source_pos_w is not None: + self.set_source_pos_w(source_pos_w) + if source_quat_w is not None: + self.set_source_quat_w(source_quat_w) + if target_pos_w is not None: + self.set_target_pos_w(target_pos_w) + if target_quat_w is not None: + self.set_target_quat_w(target_quat_w) + if target_pos_source is not None: + self.set_target_pos_source(target_pos_source) + if target_quat_source is not None: + self.set_target_quat_source(target_quat_source) + + +class MockFrameTransformer: + """Mock frame transformer sensor for testing without Isaac Sim. + + This class mimics the interface of BaseFrameTransformer for testing purposes. + It provides the same properties and methods but without simulation dependencies. + """ + + def __init__( + self, + num_instances: int, + num_target_frames: int, + target_frame_names: list[str] | None = None, + device: str = "cpu", + ): + """Initialize mock frame transformer sensor. + + Args: + num_instances: Number of environment instances. + num_target_frames: Number of target frames to track. + target_frame_names: Names of target frames. + device: Device for tensor allocation. + """ + self._num_instances = num_instances + self._num_target_frames = num_target_frames + self._target_frame_names = target_frame_names or [f"frame_{i}" for i in range(num_target_frames)] + self._device = device + self._data = MockFrameTransformerData(num_instances, num_target_frames, self._target_frame_names, device) + + # -- Properties -- + + @property + def data(self) -> MockFrameTransformerData: + """Data container for the sensor.""" + return self._data + + @property + def num_instances(self) -> int: + """Number of sensor instances.""" + return self._num_instances + + @property + def num_bodies(self) -> int: + """Number of target bodies being tracked.""" + return self._num_target_frames + + @property + def body_names(self) -> list[str]: + """Names of target bodies being tracked.""" + return self._target_frame_names + + @property + def device(self) -> str: + """Device for tensor allocation.""" + return self._device + + # -- Methods -- + + def find_bodies(self, name_keys: str | Sequence[str], preserve_order: bool = False) -> tuple[list[int], list[str]]: + """Find target frames by name regex patterns. + + Args: + name_keys: Regex pattern(s) to match frame names. + preserve_order: If True, preserve order of name_keys in output. + + Returns: + Tuple of (frame_indices, frame_names) matching the patterns. + """ + if isinstance(name_keys, str): + name_keys = [name_keys] + + matched_indices = [] + matched_names = [] + + if preserve_order: + for key in name_keys: + pattern = re.compile(key) + for i, name in enumerate(self._target_frame_names): + if pattern.fullmatch(name) and i not in matched_indices: + matched_indices.append(i) + matched_names.append(name) + else: + for i, name in enumerate(self._target_frame_names): + for key in name_keys: + pattern = re.compile(key) + if pattern.fullmatch(name): + matched_indices.append(i) + matched_names.append(name) + break + + return matched_indices, matched_names + + def reset(self, env_ids: Sequence[int] | None = None) -> None: + """Reset sensor state for specified environments. + + Args: + env_ids: Environment indices to reset. If None, resets all. + """ + # No-op for mock - data persists until explicitly changed + pass + + def update(self, dt: float, force_recompute: bool = False) -> None: + """Update sensor. + + Args: + dt: Time step since last update. + force_recompute: Force recomputation of buffers. + """ + # No-op for mock - data is set explicitly + pass diff --git a/source/isaaclab/isaaclab/test/mock_interfaces/sensors/mock_imu.py b/source/isaaclab/isaaclab/test/mock_interfaces/sensors/mock_imu.py new file mode 100644 index 00000000000..12883cfa789 --- /dev/null +++ b/source/isaaclab/isaaclab/test/mock_interfaces/sensors/mock_imu.py @@ -0,0 +1,228 @@ +# 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 + +"""Mock IMU sensor for testing without Isaac Sim.""" + +from __future__ import annotations + +from collections.abc import Sequence + +import torch + + +class MockImuData: + """Mock data container for IMU sensor. + + This class mimics the interface of BaseImuData for testing purposes. + All tensor properties return zero tensors with correct shapes if not explicitly set. + """ + + def __init__(self, num_instances: int, device: str = "cpu"): + """Initialize mock IMU data. + + Args: + num_instances: Number of sensor instances. + device: Device for tensor allocation. + """ + self._num_instances = num_instances + self._device = device + + # Internal storage for mock data + self._pos_w: torch.Tensor | None = None + self._quat_w: torch.Tensor | None = None + self._projected_gravity_b: torch.Tensor | None = None + self._lin_vel_b: torch.Tensor | None = None + self._ang_vel_b: torch.Tensor | None = None + self._lin_acc_b: torch.Tensor | None = None + self._ang_acc_b: torch.Tensor | None = None + + # -- Properties -- + + @property + def pos_w(self) -> torch.Tensor: + """Position of sensor origin in world frame. Shape: (N, 3).""" + if self._pos_w is None: + return torch.zeros(self._num_instances, 3, device=self._device) + return self._pos_w + + @property + def quat_w(self) -> torch.Tensor: + """Orientation (w, x, y, z) in world frame. Shape: (N, 4).""" + if self._quat_w is None: + # Default to identity quaternion (w, x, y, z) = (1, 0, 0, 0) + quat = torch.zeros(self._num_instances, 4, device=self._device) + quat[:, 0] = 1.0 + return quat + return self._quat_w + + @property + def pose_w(self) -> torch.Tensor: + """Pose in world frame (pos + quat). Shape: (N, 7).""" + return torch.cat([self.pos_w, self.quat_w], dim=-1) + + @property + def projected_gravity_b(self) -> torch.Tensor: + """Gravity direction in IMU body frame. Shape: (N, 3).""" + if self._projected_gravity_b is None: + # Default gravity pointing down in body frame + gravity = torch.zeros(self._num_instances, 3, device=self._device) + gravity[:, 2] = -1.0 + return gravity + return self._projected_gravity_b + + @property + def lin_vel_b(self) -> torch.Tensor: + """Linear velocity in IMU body frame. Shape: (N, 3).""" + if self._lin_vel_b is None: + return torch.zeros(self._num_instances, 3, device=self._device) + return self._lin_vel_b + + @property + def ang_vel_b(self) -> torch.Tensor: + """Angular velocity in IMU body frame. Shape: (N, 3).""" + if self._ang_vel_b is None: + return torch.zeros(self._num_instances, 3, device=self._device) + return self._ang_vel_b + + @property + def lin_acc_b(self) -> torch.Tensor: + """Linear acceleration in IMU body frame. Shape: (N, 3).""" + if self._lin_acc_b is None: + return torch.zeros(self._num_instances, 3, device=self._device) + return self._lin_acc_b + + @property + def ang_acc_b(self) -> torch.Tensor: + """Angular acceleration in IMU body frame. Shape: (N, 3).""" + if self._ang_acc_b is None: + return torch.zeros(self._num_instances, 3, device=self._device) + return self._ang_acc_b + + # -- Setters -- + + def set_pos_w(self, value: torch.Tensor) -> None: + """Set position in world frame.""" + self._pos_w = value.to(self._device) + + def set_quat_w(self, value: torch.Tensor) -> None: + """Set orientation in world frame.""" + self._quat_w = value.to(self._device) + + def set_projected_gravity_b(self, value: torch.Tensor) -> None: + """Set projected gravity in body frame.""" + self._projected_gravity_b = value.to(self._device) + + def set_lin_vel_b(self, value: torch.Tensor) -> None: + """Set linear velocity in body frame.""" + self._lin_vel_b = value.to(self._device) + + def set_ang_vel_b(self, value: torch.Tensor) -> None: + """Set angular velocity in body frame.""" + self._ang_vel_b = value.to(self._device) + + def set_lin_acc_b(self, value: torch.Tensor) -> None: + """Set linear acceleration in body frame.""" + self._lin_acc_b = value.to(self._device) + + def set_ang_acc_b(self, value: torch.Tensor) -> None: + """Set angular acceleration in body frame.""" + self._ang_acc_b = value.to(self._device) + + def set_mock_data( + self, + pos_w: torch.Tensor | None = None, + quat_w: torch.Tensor | None = None, + projected_gravity_b: torch.Tensor | None = None, + lin_vel_b: torch.Tensor | None = None, + ang_vel_b: torch.Tensor | None = None, + lin_acc_b: torch.Tensor | None = None, + ang_acc_b: torch.Tensor | None = None, + ) -> None: + """Bulk setter for mock data. + + Args: + pos_w: Position in world frame. Shape: (N, 3). + quat_w: Orientation (w, x, y, z) in world frame. Shape: (N, 4). + projected_gravity_b: Gravity direction in body frame. Shape: (N, 3). + lin_vel_b: Linear velocity in body frame. Shape: (N, 3). + ang_vel_b: Angular velocity in body frame. Shape: (N, 3). + lin_acc_b: Linear acceleration in body frame. Shape: (N, 3). + ang_acc_b: Angular acceleration in body frame. Shape: (N, 3). + """ + if pos_w is not None: + self.set_pos_w(pos_w) + if quat_w is not None: + self.set_quat_w(quat_w) + if projected_gravity_b is not None: + self.set_projected_gravity_b(projected_gravity_b) + if lin_vel_b is not None: + self.set_lin_vel_b(lin_vel_b) + if ang_vel_b is not None: + self.set_ang_vel_b(ang_vel_b) + if lin_acc_b is not None: + self.set_lin_acc_b(lin_acc_b) + if ang_acc_b is not None: + self.set_ang_acc_b(ang_acc_b) + + +class MockImu: + """Mock IMU sensor for testing without Isaac Sim. + + This class mimics the interface of BaseImu for testing purposes. + It provides the same properties and methods but without simulation dependencies. + """ + + def __init__( + self, + num_instances: int, + device: str = "cpu", + ): + """Initialize mock IMU sensor. + + Args: + num_instances: Number of sensor instances. + device: Device for tensor allocation. + """ + self._num_instances = num_instances + self._device = device + self._data = MockImuData(num_instances, device) + + # -- Properties -- + + @property + def data(self) -> MockImuData: + """Data container for the sensor.""" + return self._data + + @property + def num_instances(self) -> int: + """Number of sensor instances.""" + return self._num_instances + + @property + def device(self) -> str: + """Device for tensor allocation.""" + return self._device + + # -- Methods -- + + def reset(self, env_ids: Sequence[int] | None = None) -> None: + """Reset sensor state for specified environments. + + Args: + env_ids: Environment indices to reset. If None, resets all. + """ + # No-op for mock - data persists until explicitly changed + pass + + def update(self, dt: float, force_recompute: bool = False) -> None: + """Update sensor. + + Args: + dt: Time step since last update. + force_recompute: Force recomputation of buffers. + """ + # No-op for mock - data is set explicitly + pass diff --git a/source/isaaclab/isaaclab/test/mock_interfaces/utils/__init__.py b/source/isaaclab/isaaclab/test/mock_interfaces/utils/__init__.py new file mode 100644 index 00000000000..5935e60c1d8 --- /dev/null +++ b/source/isaaclab/isaaclab/test/mock_interfaces/utils/__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 + +"""Utilities for creating and using mock interfaces.""" + +from .mock_generator import MockArticulationBuilder, MockSensorBuilder +from .mock_wrench_composer import MockWrenchComposer +from .patching import patch_articulation, patch_sensor diff --git a/source/isaaclab/isaaclab/test/mock_interfaces/utils/mock_generator.py b/source/isaaclab/isaaclab/test/mock_interfaces/utils/mock_generator.py new file mode 100644 index 00000000000..fb139fe1e1b --- /dev/null +++ b/source/isaaclab/isaaclab/test/mock_interfaces/utils/mock_generator.py @@ -0,0 +1,350 @@ +# 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 + +"""Utilities for generating custom mock objects.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch + +if TYPE_CHECKING: + from ..assets import MockArticulation + from ..sensors import MockContactSensor, MockFrameTransformer, MockImu + + +class MockArticulationBuilder: + """Builder class for creating custom MockArticulation instances. + + Example: + >>> robot = ( + ... MockArticulationBuilder() + ... .with_joints(["hip", "knee", "ankle"], default_pos=[0.0, 0.5, -0.5]) + ... .with_bodies(["base", "thigh", "shin", "foot"]) + ... .with_fixed_base(True) + ... .with_num_instances(8) + ... .build() + ... ) + """ + + def __init__(self): + """Initialize the builder with default values.""" + self._num_instances = 1 + self._joint_names: list[str] = [] + self._body_names: list[str] = [] + self._is_fixed_base = False + self._device = "cpu" + self._default_joint_pos: list[float] | None = None + self._default_joint_vel: list[float] | None = None + self._joint_pos_limits: tuple[float, float] | None = None + self._num_fixed_tendons = 0 + self._num_spatial_tendons = 0 + self._fixed_tendon_names: list[str] | None = None + self._spatial_tendon_names: list[str] | None = None + + def with_num_instances(self, num_instances: int) -> MockArticulationBuilder: + """Set the number of articulation instances. + + Args: + num_instances: Number of instances. + + Returns: + Self for method chaining. + """ + self._num_instances = num_instances + return self + + def with_joints( + self, + joint_names: list[str], + default_pos: list[float] | None = None, + default_vel: list[float] | None = None, + ) -> MockArticulationBuilder: + """Set joint configuration. + + Args: + joint_names: Names of the joints. + default_pos: Default joint positions. + default_vel: Default joint velocities. + + Returns: + Self for method chaining. + """ + self._joint_names = joint_names + self._default_joint_pos = default_pos + self._default_joint_vel = default_vel + return self + + def with_bodies(self, body_names: list[str]) -> MockArticulationBuilder: + """Set body configuration. + + Args: + body_names: Names of the bodies. + + Returns: + Self for method chaining. + """ + self._body_names = body_names + return self + + def with_fixed_base(self, is_fixed: bool) -> MockArticulationBuilder: + """Set whether the articulation has a fixed base. + + Args: + is_fixed: True for fixed base, False for floating base. + + Returns: + Self for method chaining. + """ + self._is_fixed_base = is_fixed + return self + + def with_device(self, device: str) -> MockArticulationBuilder: + """Set the device for tensor allocation. + + Args: + device: Device string (e.g., "cpu", "cuda:0"). + + Returns: + Self for method chaining. + """ + self._device = device + return self + + def with_joint_limits(self, lower: float, upper: float) -> MockArticulationBuilder: + """Set uniform joint position limits for all joints. + + Args: + lower: Lower joint limit. + upper: Upper joint limit. + + Returns: + Self for method chaining. + """ + self._joint_pos_limits = (lower, upper) + return self + + def with_fixed_tendons(self, tendon_names: list[str]) -> MockArticulationBuilder: + """Set fixed tendon configuration. + + Args: + tendon_names: Names of fixed tendons. + + Returns: + Self for method chaining. + """ + self._fixed_tendon_names = tendon_names + self._num_fixed_tendons = len(tendon_names) + return self + + def with_spatial_tendons(self, tendon_names: list[str]) -> MockArticulationBuilder: + """Set spatial tendon configuration. + + Args: + tendon_names: Names of spatial tendons. + + Returns: + Self for method chaining. + """ + self._spatial_tendon_names = tendon_names + self._num_spatial_tendons = len(tendon_names) + return self + + def build(self) -> MockArticulation: + """Build the MockArticulation instance. + + Returns: + Configured MockArticulation instance. + """ + from ..assets import MockArticulation + + num_joints = len(self._joint_names) if self._joint_names else 1 + num_bodies = len(self._body_names) if self._body_names else num_joints + 1 + + robot = MockArticulation( + num_instances=self._num_instances, + num_joints=num_joints, + num_bodies=num_bodies, + joint_names=self._joint_names or None, + body_names=self._body_names or None, + is_fixed_base=self._is_fixed_base, + num_fixed_tendons=self._num_fixed_tendons, + num_spatial_tendons=self._num_spatial_tendons, + fixed_tendon_names=self._fixed_tendon_names, + spatial_tendon_names=self._spatial_tendon_names, + device=self._device, + ) + + # Set default joint positions + if self._default_joint_pos is not None: + default_pos = torch.tensor( + [self._default_joint_pos] * self._num_instances, + device=self._device, + ) + robot.data.set_default_joint_pos(default_pos) + robot.data.set_joint_pos(default_pos) + + # Set default joint velocities + if self._default_joint_vel is not None: + default_vel = torch.tensor( + [self._default_joint_vel] * self._num_instances, + device=self._device, + ) + robot.data.set_default_joint_vel(default_vel) + robot.data.set_joint_vel(default_vel) + + # Set joint limits + if self._joint_pos_limits is not None: + limits = torch.zeros(self._num_instances, num_joints, 2, device=self._device) + limits[..., 0] = self._joint_pos_limits[0] + limits[..., 1] = self._joint_pos_limits[1] + robot.data.set_joint_pos_limits(limits) + + return robot + + +class MockSensorBuilder: + """Builder class for creating custom mock sensor instances. + + Example: + >>> sensor = ( + ... MockSensorBuilder("contact") + ... .with_num_instances(4) + ... .with_bodies(["FL_foot", "FR_foot", "RL_foot", "RR_foot"]) + ... .with_device("cuda") + ... .build() + ... ) + """ + + def __init__(self, sensor_type: str): + """Initialize the builder. + + Args: + sensor_type: Type of sensor ("contact", "imu", or "frame_transformer"). + """ + if sensor_type not in ("contact", "imu", "frame_transformer"): + raise ValueError(f"Unknown sensor type: {sensor_type}") + self._sensor_type = sensor_type + self._num_instances = 1 + self._device = "cpu" + + # Contact sensor specific + self._body_names: list[str] = [] + self._history_length = 0 + self._num_filter_bodies = 0 + + # Frame transformer specific + self._target_frame_names: list[str] = [] + + def with_num_instances(self, num_instances: int) -> MockSensorBuilder: + """Set the number of sensor instances. + + Args: + num_instances: Number of instances. + + Returns: + Self for method chaining. + """ + self._num_instances = num_instances + return self + + def with_device(self, device: str) -> MockSensorBuilder: + """Set the device for tensor allocation. + + Args: + device: Device string. + + Returns: + Self for method chaining. + """ + self._device = device + return self + + def with_bodies(self, body_names: list[str]) -> MockSensorBuilder: + """Set body names (for contact sensor). + + Args: + body_names: Names of bodies with contact sensors. + + Returns: + Self for method chaining. + """ + self._body_names = body_names + return self + + def with_history_length(self, length: int) -> MockSensorBuilder: + """Set history buffer length (for contact sensor). + + Args: + length: Length of history buffer. + + Returns: + Self for method chaining. + """ + self._history_length = length + return self + + def with_filter_bodies(self, num_filter_bodies: int) -> MockSensorBuilder: + """Set number of filter bodies (for contact sensor). + + Args: + num_filter_bodies: Number of filter bodies for force matrix. + + Returns: + Self for method chaining. + """ + self._num_filter_bodies = num_filter_bodies + return self + + def with_target_frames(self, frame_names: list[str]) -> MockSensorBuilder: + """Set target frame names (for frame transformer). + + Args: + frame_names: Names of target frames. + + Returns: + Self for method chaining. + """ + self._target_frame_names = frame_names + return self + + def build(self) -> MockContactSensor | MockImu | MockFrameTransformer: + """Build the mock sensor instance. + + Returns: + Configured mock sensor instance. + """ + if self._sensor_type == "contact": + from ..sensors import MockContactSensor + + num_bodies = len(self._body_names) if self._body_names else 1 + return MockContactSensor( + num_instances=self._num_instances, + num_bodies=num_bodies, + body_names=self._body_names or None, + device=self._device, + history_length=self._history_length, + num_filter_bodies=self._num_filter_bodies, + ) + elif self._sensor_type == "imu": + from ..sensors import MockImu + + return MockImu( + num_instances=self._num_instances, + device=self._device, + ) + elif self._sensor_type == "frame_transformer": + from ..sensors import MockFrameTransformer + + num_frames = len(self._target_frame_names) if self._target_frame_names else 1 + return MockFrameTransformer( + num_instances=self._num_instances, + num_target_frames=num_frames, + target_frame_names=self._target_frame_names or None, + device=self._device, + ) + else: + raise ValueError(f"Unknown sensor type: {self._sensor_type}") 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 new file mode 100644 index 00000000000..4532f694420 --- /dev/null +++ b/source/isaaclab/isaaclab/test/mock_interfaces/utils/mock_wrench_composer.py @@ -0,0 +1,171 @@ +# 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 + +"""Mock WrenchComposer for testing and benchmarking. + +This module provides a mock implementation of the WrenchComposer class that can be used +in testing and benchmarking without requiring the full simulation environment. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch +import warp as wp + +if TYPE_CHECKING: + from isaaclab.assets import BaseArticulation, BaseRigidObject, BaseRigidObjectCollection + + +class MockWrenchComposer: + """Mock WrenchComposer 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. + """ + + def __init__(self, asset: BaseArticulation | BaseRigidObject | BaseRigidObjectCollection) -> None: + """Initialize the mock wrench composer. + + Args: + asset: Asset to use (Articulation, RigidObject, or RigidObjectCollection). + """ + self.num_envs = asset.num_instances + if hasattr(asset, "num_bodies"): + self.num_bodies = asset.num_bodies + else: + 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) + + # 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) + + # Create index arrays + self._ALL_ENV_INDICES_WP = wp.from_torch( + torch.arange(self.num_envs, dtype=torch.int32, device=self.device), dtype=wp.int32 + ) + self._ALL_BODY_INDICES_WP = wp.from_torch( + torch.arange(self.num_bodies, dtype=torch.int32, device=self.device), dtype=wp.int32 + ) + 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) + + @property + def active(self) -> bool: + """Whether the wrench composer is active.""" + return self._active + + @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) + """ + return self._composed_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) + """ + return self._composed_torque_b + + @property + def composed_force_as_torch(self) -> torch.Tensor: + """Composed force at the body's link frame as torch tensor. + + Returns: + torch.Tensor: Composed force at the body's link frame. (num_envs, num_bodies, 3) + """ + return self._composed_force_b_torch + + @property + def composed_torque_as_torch(self) -> torch.Tensor: + """Composed torque at the body's link frame as torch tensor. + + Returns: + torch.Tensor: Composed torque at the body's link frame. (num_envs, num_bodies, 3) + """ + return self._composed_torque_b_torch + + def add_forces_and_torques( + self, + forces: wp.array | torch.Tensor | None = None, + torques: wp.array | torch.Tensor | None = None, + positions: wp.array | torch.Tensor | None = None, + body_ids: wp.array | torch.Tensor | None = None, + env_ids: wp.array | torch.Tensor | None = None, + is_global: bool = False, + ) -> None: + """Add forces and torques (mock - just sets active flag). + + 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. 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. + """ + if forces is not None or torques is not None: + self._active = True + + def set_forces_and_torques( + self, + forces: wp.array | torch.Tensor | None = None, + torques: wp.array | torch.Tensor | None = None, + positions: wp.array | torch.Tensor | None = None, + body_ids: wp.array | torch.Tensor | None = None, + env_ids: wp.array | torch.Tensor | None = None, + is_global: bool = False, + ) -> None: + """Set forces and torques (mock - just sets active flag). + + 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. 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. + """ + if forces is not None or torques is not None: + self._active = True + + def reset(self, env_ids: wp.array | torch.Tensor | None = None) -> None: + """Reset the composed force and torque. + + Args: + env_ids: Environment ids to reset. Defaults to None (all environments). + """ + if env_ids is None: + self._composed_force_b.zero_() + self._composed_torque_b.zero_() + self._active = False + else: + # For partial reset, just zero the specified environments + 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_() diff --git a/source/isaaclab/isaaclab/test/mock_interfaces/utils/patching.py b/source/isaaclab/isaaclab/test/mock_interfaces/utils/patching.py new file mode 100644 index 00000000000..15941ce5122 --- /dev/null +++ b/source/isaaclab/isaaclab/test/mock_interfaces/utils/patching.py @@ -0,0 +1,258 @@ +# 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 + +"""Utilities for patching real classes with mock implementations in tests.""" + +from __future__ import annotations + +import functools +from collections.abc import Callable, Generator +from contextlib import contextmanager +from typing import Any, TypeVar +from unittest.mock import patch + +F = TypeVar("F", bound=Callable[..., Any]) + + +@contextmanager +def patch_articulation( + target: str, + num_instances: int = 1, + num_joints: int = 12, + num_bodies: int = 13, + joint_names: list[str] | None = None, + body_names: list[str] | None = None, + is_fixed_base: bool = False, + device: str = "cpu", + **kwargs: Any, +) -> Generator[Any, None, None]: + """Context manager for patching an articulation class with MockArticulation. + + Args: + target: The target to patch (e.g., "my_module.Articulation"). + num_instances: Number of articulation instances. + num_joints: Number of joints. + num_bodies: Number of bodies. + joint_names: Names of joints. + body_names: Names of bodies. + is_fixed_base: Whether the articulation has a fixed base. + device: Device for tensor allocation. + **kwargs: Additional keyword arguments for MockArticulation. + + Yields: + The mock articulation class. + + Example: + >>> with patch_articulation("my_module.robot", num_joints=12) as MockRobot: + ... robot = MockRobot() + ... robot.data.set_joint_pos(torch.zeros(1, 12)) + ... result = my_function_using_robot(robot) + """ + from ..assets import MockArticulation + + def create_mock(*args: Any, **create_kwargs: Any) -> MockArticulation: + # Merge configuration with any runtime kwargs + return MockArticulation( + num_instances=create_kwargs.get("num_instances", num_instances), + num_joints=create_kwargs.get("num_joints", num_joints), + num_bodies=create_kwargs.get("num_bodies", num_bodies), + joint_names=create_kwargs.get("joint_names", joint_names), + body_names=create_kwargs.get("body_names", body_names), + is_fixed_base=create_kwargs.get("is_fixed_base", is_fixed_base), + device=create_kwargs.get("device", device), + **{k: v for k, v in kwargs.items() if k not in create_kwargs}, + ) + + with patch(target, side_effect=create_mock) as mock_class: + yield mock_class + + +@contextmanager +def patch_sensor( + target: str, + sensor_type: str, + num_instances: int = 1, + device: str = "cpu", + **kwargs: Any, +) -> Generator[Any, None, None]: + """Context manager for patching a sensor class with a mock sensor. + + Args: + target: The target to patch (e.g., "my_module.ContactSensor"). + sensor_type: Type of sensor ("contact", "imu", or "frame_transformer"). + num_instances: Number of sensor instances. + device: Device for tensor allocation. + **kwargs: Additional keyword arguments for the mock sensor. + + Yields: + The mock sensor class. + + Example: + >>> with patch_sensor("my_env.ContactSensor", "contact", num_bodies=4) as MockSensor: + ... sensor = MockSensor() + ... sensor.data.set_net_forces_w(torch.randn(1, 4, 3)) + ... result = my_function(sensor) + """ + if sensor_type == "contact": + from ..sensors import MockContactSensor + + def create_mock(*args: Any, **create_kwargs: Any) -> MockContactSensor: + return MockContactSensor( + num_instances=create_kwargs.get("num_instances", num_instances), + num_bodies=create_kwargs.get("num_bodies", kwargs.get("num_bodies", 1)), + body_names=create_kwargs.get("body_names", kwargs.get("body_names")), + device=create_kwargs.get("device", device), + history_length=create_kwargs.get("history_length", kwargs.get("history_length", 0)), + num_filter_bodies=create_kwargs.get("num_filter_bodies", kwargs.get("num_filter_bodies", 0)), + ) + + elif sensor_type == "imu": + from ..sensors import MockImu + + def create_mock(*args: Any, **create_kwargs: Any) -> MockImu: + return MockImu( + num_instances=create_kwargs.get("num_instances", num_instances), + device=create_kwargs.get("device", device), + ) + + elif sensor_type == "frame_transformer": + from ..sensors import MockFrameTransformer + + def create_mock(*args: Any, **create_kwargs: Any) -> MockFrameTransformer: + return MockFrameTransformer( + num_instances=create_kwargs.get("num_instances", num_instances), + num_target_frames=create_kwargs.get("num_target_frames", kwargs.get("num_target_frames", 1)), + target_frame_names=create_kwargs.get("target_frame_names", kwargs.get("target_frame_names")), + device=create_kwargs.get("device", device), + ) + + else: + raise ValueError(f"Unknown sensor type: {sensor_type}") + + with patch(target, side_effect=create_mock) as mock_class: + yield mock_class + + +def mock_articulation( + num_instances: int = 1, + num_joints: int = 12, + num_bodies: int = 13, + joint_names: list[str] | None = None, + body_names: list[str] | None = None, + is_fixed_base: bool = False, + device: str = "cpu", + **kwargs: Any, +) -> Callable[[F], F]: + """Decorator for injecting a MockArticulation into a test function. + + The mock articulation is passed as the first argument to the decorated function. + + Args: + num_instances: Number of articulation instances. + num_joints: Number of joints. + num_bodies: Number of bodies. + joint_names: Names of joints. + body_names: Names of bodies. + is_fixed_base: Whether the articulation has a fixed base. + device: Device for tensor allocation. + **kwargs: Additional keyword arguments for MockArticulation. + + Returns: + Decorator function. + + Example: + >>> @mock_articulation(num_joints=12, num_bodies=13) + ... def test_my_function(mock_robot): + ... mock_robot.data.set_joint_pos(torch.zeros(1, 12)) + ... result = compute_something(mock_robot) + ... assert result.shape == (1, 12) + """ + from ..assets import MockArticulation + + def decorator(func: F) -> F: + @functools.wraps(func) + def wrapper(*args: Any, **wrap_kwargs: Any) -> Any: + mock = MockArticulation( + num_instances=num_instances, + num_joints=num_joints, + num_bodies=num_bodies, + joint_names=joint_names, + body_names=body_names, + is_fixed_base=is_fixed_base, + device=device, + **kwargs, + ) + return func(mock, *args, **wrap_kwargs) + + return wrapper # type: ignore + + return decorator + + +def mock_sensor( + sensor_type: str, + num_instances: int = 1, + device: str = "cpu", + **kwargs: Any, +) -> Callable[[F], F]: + """Decorator for injecting a mock sensor into a test function. + + The mock sensor is passed as the first argument to the decorated function. + + Args: + sensor_type: Type of sensor ("contact", "imu", or "frame_transformer"). + num_instances: Number of sensor instances. + device: Device for tensor allocation. + **kwargs: Additional keyword arguments for the mock sensor. + + Returns: + Decorator function. + + Example: + >>> @mock_sensor("contact", num_instances=4, num_bodies=4) + ... def test_contact_reward(mock_contact): + ... mock_contact.data.set_net_forces_w(torch.randn(4, 4, 3)) + ... reward = compute_contact_reward(mock_contact) + ... assert reward.shape == (4,) + """ + + def decorator(func: F) -> F: + @functools.wraps(func) + def wrapper(*args: Any, **wrap_kwargs: Any) -> Any: + if sensor_type == "contact": + from ..sensors import MockContactSensor + + mock = MockContactSensor( + num_instances=num_instances, + num_bodies=kwargs.get("num_bodies", 1), + body_names=kwargs.get("body_names"), + device=device, + history_length=kwargs.get("history_length", 0), + num_filter_bodies=kwargs.get("num_filter_bodies", 0), + ) + elif sensor_type == "imu": + from ..sensors import MockImu + + mock = MockImu( + num_instances=num_instances, + device=device, + ) + elif sensor_type == "frame_transformer": + from ..sensors import MockFrameTransformer + + mock = MockFrameTransformer( + num_instances=num_instances, + num_target_frames=kwargs.get("num_target_frames", 1), + target_frame_names=kwargs.get("target_frame_names"), + device=device, + ) + else: + raise ValueError(f"Unknown sensor type: {sensor_type}") + + return func(mock, *args, **wrap_kwargs) + + return wrapper # type: ignore + + return decorator diff --git a/source/isaaclab/test/test_mock_interfaces/test_mock_assets.py b/source/isaaclab/test/test_mock_interfaces/test_mock_assets.py new file mode 100644 index 00000000000..37e5b05914b --- /dev/null +++ b/source/isaaclab/test/test_mock_interfaces/test_mock_assets.py @@ -0,0 +1,335 @@ +# 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 mock asset interfaces.""" + +import pytest +import torch + +from isaaclab.test.mock_interfaces.assets import ( + MockArticulation, + MockRigidObject, + MockRigidObjectCollection, + create_mock_articulation, + create_mock_humanoid, + create_mock_quadruped, + create_mock_rigid_object, + create_mock_rigid_object_collection, +) +from isaaclab.test.mock_interfaces.utils import MockArticulationBuilder + +# ============================================================================== +# MockArticulation Tests +# ============================================================================== + + +class TestMockArticulation: + """Tests for MockArticulation and MockArticulationData.""" + + @pytest.fixture + def robot(self): + """Create a mock articulation fixture.""" + return MockArticulation( + num_instances=4, + num_joints=12, + num_bodies=13, + device="cpu", + ) + + def test_initialization(self, robot): + """Test that MockArticulation initializes correctly.""" + assert robot.num_instances == 4 + assert robot.num_joints == 12 + assert robot.num_bodies == 13 + assert robot.device == "cpu" + assert robot.root_view is None + assert robot.data is not None + + def test_joint_state_shapes(self, robot): + """Test joint state tensor shapes.""" + assert robot.data.joint_pos.shape == (4, 12) + assert robot.data.joint_vel.shape == (4, 12) + assert robot.data.joint_acc.shape == (4, 12) + + def test_root_state_shapes(self, robot): + """Test root state tensor shapes.""" + # Link frame + assert robot.data.root_link_pose_w.shape == (4, 7) + assert robot.data.root_link_vel_w.shape == (4, 6) + assert robot.data.root_link_state_w.shape == (4, 13) + + # Sliced properties + assert robot.data.root_link_pos_w.shape == (4, 3) + assert robot.data.root_link_quat_w.shape == (4, 4) + assert robot.data.root_link_lin_vel_w.shape == (4, 3) + assert robot.data.root_link_ang_vel_w.shape == (4, 3) + + def test_body_state_shapes(self, robot): + """Test body state tensor shapes.""" + assert robot.data.body_link_pose_w.shape == (4, 13, 7) + assert robot.data.body_link_vel_w.shape == (4, 13, 6) + assert robot.data.body_link_state_w.shape == (4, 13, 13) + + def test_default_state_shapes(self, robot): + """Test default state tensor shapes.""" + assert robot.data.default_root_pose.shape == (4, 7) + assert robot.data.default_root_vel.shape == (4, 6) + assert robot.data.default_root_state.shape == (4, 13) + assert robot.data.default_joint_pos.shape == (4, 12) + assert robot.data.default_joint_vel.shape == (4, 12) + + def test_identity_quaternion_default(self, robot): + """Test that default quaternions are identity quaternions.""" + quat = robot.data.root_link_quat_w + # w=1, x=y=z=0 + assert torch.all(quat[:, 0] == 1) + assert torch.all(quat[:, 1:] == 0) + + def test_set_joint_pos(self, robot): + """Test setting joint positions.""" + joint_pos = torch.randn(4, 12) + robot.data.set_joint_pos(joint_pos) + assert torch.allclose(robot.data.joint_pos, joint_pos) + + def test_set_mock_data_bulk(self, robot): + """Test bulk data setter.""" + joint_pos = torch.randn(4, 12) + joint_vel = torch.randn(4, 12) + + robot.data.set_mock_data(joint_pos=joint_pos, joint_vel=joint_vel) + + assert torch.allclose(robot.data.joint_pos, joint_pos) + assert torch.allclose(robot.data.joint_vel, joint_vel) + + def test_find_joints(self): + """Test joint finding by regex.""" + joint_names = ["FL_hip", "FL_thigh", "FL_calf", "FR_hip", "FR_thigh", "FR_calf"] + robot = MockArticulation( + num_instances=1, + num_joints=6, + num_bodies=7, + joint_names=joint_names, + device="cpu", + ) + + # Find all hip joints + indices, names = robot.find_joints(".*_hip") + assert len(indices) == 2 + assert "FL_hip" in names + assert "FR_hip" in names + + # Find FL leg joints + indices, names = robot.find_joints("FL_.*") + assert len(indices) == 3 + + def test_find_bodies(self): + """Test body finding by regex.""" + body_names = ["base", "FL_hip", "FL_thigh", "FL_calf", "FR_hip", "FR_thigh", "FR_calf"] + robot = MockArticulation( + num_instances=1, + num_joints=6, + num_bodies=7, + body_names=body_names, + device="cpu", + ) + + # Find base + indices, names = robot.find_bodies("base") + assert indices == [0] + + # Find all thigh bodies + indices, names = robot.find_bodies(".*_thigh") + assert len(indices) == 2 + + def test_set_joint_position_target(self, robot): + """Test setting joint position targets.""" + target = torch.randn(4, 12) + robot.set_joint_position_target(target) + assert torch.allclose(robot.data.joint_pos_target, target) + + def test_joint_limits(self, robot): + """Test joint limits.""" + limits = robot.data.joint_pos_limits + assert limits.shape == (4, 12, 2) + # Default limits should be -inf to inf + assert torch.all(limits[..., 0] == float("-inf")) + assert torch.all(limits[..., 1] == float("inf")) + + +# ============================================================================== +# MockRigidObject Tests +# ============================================================================== + + +class TestMockRigidObject: + """Tests for MockRigidObject and MockRigidObjectData.""" + + @pytest.fixture + def obj(self): + """Create a mock rigid object fixture.""" + return MockRigidObject(num_instances=4, device="cpu") + + def test_initialization(self, obj): + """Test that MockRigidObject initializes correctly.""" + assert obj.num_instances == 4 + assert obj.num_bodies == 1 # Always 1 for rigid object + assert obj.root_view is None + + def test_root_state_shapes(self, obj): + """Test root state tensor shapes.""" + assert obj.data.root_link_pose_w.shape == (4, 7) + assert obj.data.root_link_vel_w.shape == (4, 6) + assert obj.data.root_link_state_w.shape == (4, 13) + + def test_body_state_shapes(self, obj): + """Test body state tensor shapes (single body).""" + assert obj.data.body_link_pose_w.shape == (4, 1, 7) + assert obj.data.body_link_vel_w.shape == (4, 1, 6) + + def test_body_properties(self, obj): + """Test body property shapes.""" + assert obj.data.body_mass.shape == (4, 1, 1) + assert obj.data.body_inertia.shape == (4, 1, 9) + + +# ============================================================================== +# MockRigidObjectCollection Tests +# ============================================================================== + + +class TestMockRigidObjectCollection: + """Tests for MockRigidObjectCollection and MockRigidObjectCollectionData.""" + + @pytest.fixture + def collection(self): + """Create a mock rigid object collection fixture.""" + return MockRigidObjectCollection( + num_instances=4, + num_bodies=5, + device="cpu", + ) + + def test_initialization(self, collection): + """Test that MockRigidObjectCollection initializes correctly.""" + assert collection.num_instances == 4 + assert collection.num_bodies == 5 + + def test_body_state_shapes(self, collection): + """Test body state tensor shapes.""" + assert collection.data.body_link_pose_w.shape == (4, 5, 7) + assert collection.data.body_link_vel_w.shape == (4, 5, 6) + assert collection.data.body_link_state_w.shape == (4, 5, 13) + + def test_find_bodies_returns_mask(self, collection): + """Test that find_bodies returns a mask tensor.""" + body_mask, names, indices = collection.find_bodies("body_0") + assert isinstance(body_mask, torch.Tensor) + assert body_mask.shape == (5,) + assert body_mask[0] + + +# ============================================================================== +# Factory Function Tests +# ============================================================================== + + +class TestAssetFactories: + """Tests for asset factory functions.""" + + def test_create_mock_quadruped(self): + """Test quadruped factory function.""" + robot = create_mock_quadruped(num_instances=4) + assert robot.num_instances == 4 + assert robot.num_joints == 12 + assert robot.num_bodies == 13 + assert "FL_hip" in robot.joint_names + assert "base" in robot.body_names + + def test_create_mock_humanoid(self): + """Test humanoid factory function.""" + robot = create_mock_humanoid(num_instances=2) + assert robot.num_instances == 2 + assert robot.num_joints == 21 + + def test_create_mock_articulation(self): + """Test generic articulation factory function.""" + robot = create_mock_articulation( + num_instances=2, + num_joints=6, + num_bodies=7, + is_fixed_base=True, + ) + assert robot.num_instances == 2 + assert robot.num_joints == 6 + assert robot.is_fixed_base + + def test_create_mock_rigid_object(self): + """Test rigid object factory function.""" + obj = create_mock_rigid_object(num_instances=3) + assert obj.num_instances == 3 + assert obj.num_bodies == 1 + assert obj.data.root_link_pose_w.shape == (3, 7) + + def test_create_mock_rigid_object_collection(self): + """Test rigid object collection factory function.""" + collection = create_mock_rigid_object_collection( + num_instances=4, + num_bodies=6, + body_names=["obj_0", "obj_1", "obj_2", "obj_3", "obj_4", "obj_5"], + ) + assert collection.num_instances == 4 + assert collection.num_bodies == 6 + assert collection.body_names == ["obj_0", "obj_1", "obj_2", "obj_3", "obj_4", "obj_5"] + assert collection.data.body_link_pose_w.shape == (4, 6, 7) + + +# ============================================================================== +# MockArticulationBuilder Tests +# ============================================================================== + + +class TestMockArticulationBuilder: + """Tests for MockArticulationBuilder.""" + + def test_basic_build(self): + """Test building a basic articulation.""" + robot = ( + MockArticulationBuilder() + .with_num_instances(4) + .with_joints(["joint_0", "joint_1", "joint_2"]) + .with_bodies(["base", "link_1", "link_2", "link_3"]) + .build() + ) + + assert robot.num_instances == 4 + assert robot.num_joints == 3 + assert robot.num_bodies == 4 + + def test_with_default_positions(self): + """Test setting default joint positions.""" + default_pos = [0.0, 0.5, -0.5] + robot = ( + MockArticulationBuilder() + .with_num_instances(2) + .with_joints(["j0", "j1", "j2"], default_pos=default_pos) + .build() + ) + + expected = torch.tensor([default_pos, default_pos]) + assert torch.allclose(robot.data.joint_pos, expected) + + def test_with_joint_limits(self): + """Test setting joint limits.""" + robot = ( + MockArticulationBuilder() + .with_num_instances(1) + .with_joints(["j0", "j1"]) + .with_joint_limits(-1.0, 1.0) + .build() + ) + + limits = robot.data.joint_pos_limits + assert torch.all(limits[..., 0] == -1.0) + assert torch.all(limits[..., 1] == 1.0) diff --git a/source/isaaclab/test/test_mock_interfaces/test_mock_data_properties.py b/source/isaaclab/test/test_mock_interfaces/test_mock_data_properties.py new file mode 100644 index 00000000000..37e09e2f382 --- /dev/null +++ b/source/isaaclab/test/test_mock_interfaces/test_mock_data_properties.py @@ -0,0 +1,835 @@ +# 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 + +"""Comprehensive tests for mock data properties - shapes, setters, and device handling.""" + +import pytest +import torch + +from isaaclab.test.mock_interfaces.assets import ( + MockArticulationData, + MockRigidObjectCollectionData, + MockRigidObjectData, +) +from isaaclab.test.mock_interfaces.sensors import ( + MockContactSensorData, + MockFrameTransformerData, + MockImuData, +) + +# ============================================================================== +# IMU Data Property Tests +# ============================================================================== + + +class TestMockImuDataProperties: + """Comprehensive tests for all MockImuData properties.""" + + @pytest.fixture + def data(self): + """Create MockImuData fixture.""" + return MockImuData(num_instances=4, device="cpu") + + @pytest.mark.parametrize( + "property_name,expected_shape", + [ + ("pos_w", (4, 3)), + ("quat_w", (4, 4)), + ("pose_w", (4, 7)), + ("projected_gravity_b", (4, 3)), + ("lin_vel_b", (4, 3)), + ("ang_vel_b", (4, 3)), + ("lin_acc_b", (4, 3)), + ("ang_acc_b", (4, 3)), + ], + ) + def test_property_shapes(self, data, property_name, expected_shape): + """Test that all properties return tensors with correct shapes.""" + prop = getattr(data, property_name) + assert prop.shape == expected_shape, f"{property_name} has wrong shape" + + @pytest.mark.parametrize( + "setter_name,property_name,shape", + [ + ("set_pos_w", "pos_w", (4, 3)), + ("set_quat_w", "quat_w", (4, 4)), + ("set_projected_gravity_b", "projected_gravity_b", (4, 3)), + ("set_lin_vel_b", "lin_vel_b", (4, 3)), + ("set_ang_vel_b", "ang_vel_b", (4, 3)), + ("set_lin_acc_b", "lin_acc_b", (4, 3)), + ("set_ang_acc_b", "ang_acc_b", (4, 3)), + ], + ) + def test_setters_update_properties(self, data, setter_name, property_name, shape): + """Test that setters properly update the corresponding properties.""" + test_value = torch.randn(shape) + setter = getattr(data, setter_name) + setter(test_value) + result = getattr(data, property_name) + assert torch.allclose(result, test_value), f"{setter_name} did not update {property_name}" + + def test_bulk_setter(self, data): + """Test that set_mock_data updates multiple properties at once.""" + lin_vel = torch.randn(4, 3) + ang_vel = torch.randn(4, 3) + lin_acc = torch.randn(4, 3) + + data.set_mock_data(lin_vel_b=lin_vel, ang_vel_b=ang_vel, lin_acc_b=lin_acc) + + assert torch.allclose(data.lin_vel_b, lin_vel) + assert torch.allclose(data.ang_vel_b, ang_vel) + assert torch.allclose(data.lin_acc_b, lin_acc) + + def test_default_quaternion_is_identity(self, data): + """Test that default quaternion is identity (w=1, x=y=z=0).""" + quat = data.quat_w + assert torch.allclose(quat[:, 0], torch.ones(4)) # w=1 + assert torch.allclose(quat[:, 1:], torch.zeros(4, 3)) # xyz=0 + + def test_default_gravity_points_down(self, data): + """Test that default gravity points in -z direction.""" + gravity = data.projected_gravity_b + expected = torch.tensor([[0, 0, -1]] * 4, dtype=torch.float32) + assert torch.allclose(gravity, expected) + + +# ============================================================================== +# Contact Sensor Data Property Tests +# ============================================================================== + + +class TestMockContactSensorDataProperties: + """Comprehensive tests for all MockContactSensorData properties.""" + + @pytest.fixture + def data(self): + """Create MockContactSensorData fixture.""" + return MockContactSensorData( + num_instances=4, + num_bodies=3, + device="cpu", + history_length=2, + num_filter_bodies=5, + ) + + @pytest.fixture + def data_no_history(self): + """Create MockContactSensorData without history.""" + return MockContactSensorData( + num_instances=4, + num_bodies=3, + device="cpu", + history_length=0, + num_filter_bodies=0, + ) + + @pytest.mark.parametrize( + "property_name,expected_shape", + [ + ("pos_w", (4, 3, 3)), + ("quat_w", (4, 3, 4)), + ("pose_w", (4, 3, 7)), + ("net_forces_w", (4, 3, 3)), + ("last_air_time", (4, 3)), + ("current_air_time", (4, 3)), + ("last_contact_time", (4, 3)), + ("current_contact_time", (4, 3)), + ], + ) + def test_basic_property_shapes(self, data, property_name, expected_shape): + """Test basic properties return tensors with correct shapes.""" + prop = getattr(data, property_name) + assert prop.shape == expected_shape, f"{property_name} has wrong shape" + + @pytest.mark.parametrize( + "property_name,expected_shape", + [ + ("net_forces_w_history", (4, 2, 3, 3)), # (N, T, B, 3) + ("force_matrix_w", (4, 3, 5, 3)), # (N, B, M, 3) + ("force_matrix_w_history", (4, 2, 3, 5, 3)), # (N, T, B, M, 3) + ("contact_pos_w", (4, 3, 5, 3)), # (N, B, M, 3) + ("friction_forces_w", (4, 3, 5, 3)), # (N, B, M, 3) + ], + ) + def test_optional_property_shapes_with_history(self, data, property_name, expected_shape): + """Test optional properties with history/filter enabled.""" + prop = getattr(data, property_name) + assert prop is not None + assert prop.shape == expected_shape, f"{property_name} has wrong shape" + + def test_optional_properties_none_without_config(self, data_no_history): + """Test optional properties are None when not configured.""" + assert data_no_history.net_forces_w_history is None + assert data_no_history.force_matrix_w is None + assert data_no_history.force_matrix_w_history is None + assert data_no_history.contact_pos_w is None + assert data_no_history.friction_forces_w is None + + @pytest.mark.parametrize( + "setter_name,property_name,shape", + [ + ("set_pos_w", "pos_w", (4, 3, 3)), + ("set_quat_w", "quat_w", (4, 3, 4)), + ("set_net_forces_w", "net_forces_w", (4, 3, 3)), + ("set_last_air_time", "last_air_time", (4, 3)), + ("set_current_air_time", "current_air_time", (4, 3)), + ("set_last_contact_time", "last_contact_time", (4, 3)), + ("set_current_contact_time", "current_contact_time", (4, 3)), + ], + ) + def test_setters_update_properties(self, data, setter_name, property_name, shape): + """Test that setters properly update the corresponding properties.""" + test_value = torch.randn(shape) + setter = getattr(data, setter_name) + setter(test_value) + result = getattr(data, property_name) + assert torch.allclose(result, test_value), f"{setter_name} did not update {property_name}" + + +# ============================================================================== +# Frame Transformer Data Property Tests +# ============================================================================== + + +class TestMockFrameTransformerDataProperties: + """Comprehensive tests for all MockFrameTransformerData properties.""" + + @pytest.fixture + def data(self): + """Create MockFrameTransformerData fixture.""" + return MockFrameTransformerData( + num_instances=4, + num_target_frames=3, + target_frame_names=["frame_0", "frame_1", "frame_2"], + device="cpu", + ) + + @pytest.mark.parametrize( + "property_name,expected_shape", + [ + ("source_pos_w", (4, 3)), + ("source_quat_w", (4, 4)), + ("source_pose_w", (4, 7)), + ("target_pos_w", (4, 3, 3)), + ("target_quat_w", (4, 3, 4)), + ("target_pose_w", (4, 3, 7)), + ("target_pos_source", (4, 3, 3)), + ("target_quat_source", (4, 3, 4)), + ("target_pose_source", (4, 3, 7)), + ], + ) + def test_property_shapes(self, data, property_name, expected_shape): + """Test that all properties return tensors with correct shapes.""" + prop = getattr(data, property_name) + assert prop.shape == expected_shape, f"{property_name} has wrong shape" + + @pytest.mark.parametrize( + "setter_name,property_name,shape", + [ + ("set_source_pos_w", "source_pos_w", (4, 3)), + ("set_source_quat_w", "source_quat_w", (4, 4)), + ("set_target_pos_w", "target_pos_w", (4, 3, 3)), + ("set_target_quat_w", "target_quat_w", (4, 3, 4)), + ("set_target_pos_source", "target_pos_source", (4, 3, 3)), + ("set_target_quat_source", "target_quat_source", (4, 3, 4)), + ], + ) + def test_setters_update_properties(self, data, setter_name, property_name, shape): + """Test that setters properly update the corresponding properties.""" + test_value = torch.randn(shape) + setter = getattr(data, setter_name) + setter(test_value) + result = getattr(data, property_name) + assert torch.allclose(result, test_value), f"{setter_name} did not update {property_name}" + + def test_target_frame_names(self, data): + """Test that target frame names are correctly stored.""" + assert data.target_frame_names == ["frame_0", "frame_1", "frame_2"] + + +# ============================================================================== +# Articulation Data Property Tests +# ============================================================================== + + +class TestMockArticulationDataProperties: + """Comprehensive tests for all MockArticulationData properties.""" + + @pytest.fixture + def data(self): + """Create MockArticulationData fixture.""" + return MockArticulationData( + num_instances=4, + num_joints=12, + num_bodies=13, + num_fixed_tendons=2, + num_spatial_tendons=1, + device="cpu", + ) + + # -- Joint State Properties -- + @pytest.mark.parametrize( + "property_name,expected_shape", + [ + ("joint_pos", (4, 12)), + ("joint_vel", (4, 12)), + ("joint_acc", (4, 12)), + ("joint_pos_target", (4, 12)), + ("joint_vel_target", (4, 12)), + ("joint_effort_target", (4, 12)), + ("computed_torque", (4, 12)), + ("applied_torque", (4, 12)), + ], + ) + def test_joint_state_shapes(self, data, property_name, expected_shape): + """Test joint state properties have correct shapes.""" + prop = getattr(data, property_name) + assert prop.shape == expected_shape, f"{property_name} has wrong shape" + + # -- Joint Property Shapes -- + @pytest.mark.parametrize( + "property_name,expected_shape", + [ + ("joint_stiffness", (4, 12)), + ("joint_damping", (4, 12)), + ("joint_armature", (4, 12)), + ("joint_friction_coeff", (4, 12)), + ("joint_dynamic_friction_coeff", (4, 12)), + ("joint_viscous_friction_coeff", (4, 12)), + ("joint_pos_limits", (4, 12, 2)), + ("joint_vel_limits", (4, 12)), + ("joint_effort_limits", (4, 12)), + ("soft_joint_pos_limits", (4, 12, 2)), + ("soft_joint_vel_limits", (4, 12)), + ("gear_ratio", (4, 12)), + ], + ) + def test_joint_property_shapes(self, data, property_name, expected_shape): + """Test joint property shapes.""" + prop = getattr(data, property_name) + assert prop.shape == expected_shape, f"{property_name} has wrong shape" + + # -- Root State Properties -- + @pytest.mark.parametrize( + "property_name,expected_shape", + [ + ("root_link_pose_w", (4, 7)), + ("root_link_vel_w", (4, 6)), + ("root_link_state_w", (4, 13)), + ("root_link_pos_w", (4, 3)), + ("root_link_quat_w", (4, 4)), + ("root_link_lin_vel_w", (4, 3)), + ("root_link_ang_vel_w", (4, 3)), + ("root_com_pose_w", (4, 7)), + ("root_com_vel_w", (4, 6)), + ("root_com_state_w", (4, 13)), + ("root_com_pos_w", (4, 3)), + ("root_com_quat_w", (4, 4)), + ("root_com_lin_vel_w", (4, 3)), + ("root_com_ang_vel_w", (4, 3)), + ("root_state_w", (4, 13)), + ], + ) + def test_root_state_shapes(self, data, property_name, expected_shape): + """Test root state properties have correct shapes.""" + prop = getattr(data, property_name) + assert prop.shape == expected_shape, f"{property_name} has wrong shape" + + # -- Body State Properties -- + @pytest.mark.parametrize( + "property_name,expected_shape", + [ + ("body_link_pose_w", (4, 13, 7)), + ("body_link_vel_w", (4, 13, 6)), + ("body_link_state_w", (4, 13, 13)), + ("body_link_pos_w", (4, 13, 3)), + ("body_link_quat_w", (4, 13, 4)), + ("body_link_lin_vel_w", (4, 13, 3)), + ("body_link_ang_vel_w", (4, 13, 3)), + ("body_com_pose_w", (4, 13, 7)), + ("body_com_vel_w", (4, 13, 6)), + ("body_com_state_w", (4, 13, 13)), + ("body_com_acc_w", (4, 13, 6)), + ("body_com_pos_w", (4, 13, 3)), + ("body_com_quat_w", (4, 13, 4)), + ("body_com_lin_vel_w", (4, 13, 3)), + ("body_com_ang_vel_w", (4, 13, 3)), + ("body_com_lin_acc_w", (4, 13, 3)), + ("body_com_ang_acc_w", (4, 13, 3)), + ], + ) + def test_body_state_shapes(self, data, property_name, expected_shape): + """Test body state properties have correct shapes.""" + prop = getattr(data, property_name) + assert prop.shape == expected_shape, f"{property_name} has wrong shape" + + # -- Body Properties -- + @pytest.mark.parametrize( + "property_name,expected_shape", + [ + ("body_mass", (4, 13)), + ("body_inertia", (4, 13, 3, 3)), + ("body_incoming_joint_wrench_b", (4, 13, 6)), + ], + ) + def test_body_property_shapes(self, data, property_name, expected_shape): + """Test body property shapes.""" + prop = getattr(data, property_name) + assert prop.shape == expected_shape, f"{property_name} has wrong shape" + + # -- Default State Properties -- + @pytest.mark.parametrize( + "property_name,expected_shape", + [ + ("default_root_pose", (4, 7)), + ("default_root_vel", (4, 6)), + ("default_root_state", (4, 13)), + ("default_joint_pos", (4, 12)), + ("default_joint_vel", (4, 12)), + ], + ) + def test_default_state_shapes(self, data, property_name, expected_shape): + """Test default state properties have correct shapes.""" + prop = getattr(data, property_name) + assert prop.shape == expected_shape, f"{property_name} has wrong shape" + + # -- Derived Properties -- + @pytest.mark.parametrize( + "property_name,expected_shape", + [ + ("projected_gravity_b", (4, 3)), + ("heading_w", (4,)), + ("root_link_lin_vel_b", (4, 3)), + ("root_link_ang_vel_b", (4, 3)), + ("root_com_lin_vel_b", (4, 3)), + ("root_com_ang_vel_b", (4, 3)), + ], + ) + def test_derived_property_shapes(self, data, property_name, expected_shape): + """Test derived properties have correct shapes.""" + prop = getattr(data, property_name) + assert prop.shape == expected_shape, f"{property_name} has wrong shape" + + # -- Tendon Properties -- + @pytest.mark.parametrize( + "property_name,expected_shape", + [ + ("fixed_tendon_stiffness", (4, 2)), + ("fixed_tendon_damping", (4, 2)), + ("fixed_tendon_limit_stiffness", (4, 2)), + ("fixed_tendon_rest_length", (4, 2)), + ("fixed_tendon_offset", (4, 2)), + ("fixed_tendon_pos_limits", (4, 2, 2)), + ("spatial_tendon_stiffness", (4, 1)), + ("spatial_tendon_damping", (4, 1)), + ("spatial_tendon_limit_stiffness", (4, 1)), + ("spatial_tendon_offset", (4, 1)), + ], + ) + def test_tendon_property_shapes(self, data, property_name, expected_shape): + """Test tendon properties have correct shapes.""" + prop = getattr(data, property_name) + assert prop.shape == expected_shape, f"{property_name} has wrong shape" + + # -- Setter Tests -- + @pytest.mark.parametrize( + "setter_name,property_name,shape", + [ + ("set_joint_pos", "joint_pos", (4, 12)), + ("set_joint_vel", "joint_vel", (4, 12)), + ("set_joint_acc", "joint_acc", (4, 12)), + ("set_root_link_pose_w", "root_link_pose_w", (4, 7)), + ("set_root_link_vel_w", "root_link_vel_w", (4, 6)), + ("set_body_link_pose_w", "body_link_pose_w", (4, 13, 7)), + ("set_body_link_vel_w", "body_link_vel_w", (4, 13, 6)), + ("set_body_mass", "body_mass", (4, 13)), + ], + ) + def test_setters_update_properties(self, data, setter_name, property_name, shape): + """Test that setters properly update the corresponding properties.""" + test_value = torch.randn(shape) + setter = getattr(data, setter_name) + setter(test_value) + result = getattr(data, property_name) + assert torch.allclose(result, test_value), f"{setter_name} did not update {property_name}" + + def test_bulk_setter_with_multiple_properties(self, data): + """Test set_mock_data with multiple properties.""" + joint_pos = torch.randn(4, 12) + joint_vel = torch.randn(4, 12) + root_pose = torch.randn(4, 7) + + data.set_mock_data( + joint_pos=joint_pos, + joint_vel=joint_vel, + root_link_pose_w=root_pose, + ) + + assert torch.allclose(data.joint_pos, joint_pos) + assert torch.allclose(data.joint_vel, joint_vel) + assert torch.allclose(data.root_link_pose_w, root_pose) + + def test_bulk_setter_unknown_property_raises(self, data): + """Test that set_mock_data raises on unknown property.""" + with pytest.raises(ValueError, match="Unknown property"): + data.set_mock_data(unknown_property=torch.zeros(4)) + + +# ============================================================================== +# Rigid Object Data Property Tests +# ============================================================================== + + +class TestMockRigidObjectDataProperties: + """Comprehensive tests for MockRigidObjectData properties.""" + + @pytest.fixture + def data(self): + """Create MockRigidObjectData fixture.""" + return MockRigidObjectData(num_instances=4, device="cpu") + + @pytest.mark.parametrize( + "property_name,expected_shape", + [ + ("root_link_pose_w", (4, 7)), + ("root_link_vel_w", (4, 6)), + ("root_link_state_w", (4, 13)), + ("root_link_pos_w", (4, 3)), + ("root_link_quat_w", (4, 4)), + ("root_link_lin_vel_w", (4, 3)), + ("root_link_ang_vel_w", (4, 3)), + ("root_com_pose_w", (4, 7)), + ("root_com_vel_w", (4, 6)), + ("root_com_state_w", (4, 13)), + ("root_state_w", (4, 13)), + ("body_link_pose_w", (4, 1, 7)), + ("body_link_vel_w", (4, 1, 6)), + ("body_link_state_w", (4, 1, 13)), + ("body_com_pose_w", (4, 1, 7)), + ("body_com_vel_w", (4, 1, 6)), + ("body_com_state_w", (4, 1, 13)), + ("body_com_acc_w", (4, 1, 6)), + ("body_mass", (4, 1, 1)), + ("body_inertia", (4, 1, 9)), + ("projected_gravity_b", (4, 3)), + ("heading_w", (4,)), + ("default_root_pose", (4, 7)), + ("default_root_vel", (4, 6)), + ("default_root_state", (4, 13)), + ], + ) + def test_property_shapes(self, data, property_name, expected_shape): + """Test that all properties return tensors with correct shapes.""" + prop = getattr(data, property_name) + assert prop.shape == expected_shape, f"{property_name} has wrong shape" + + +# ============================================================================== +# Rigid Object Collection Data Property Tests +# ============================================================================== + + +class TestMockRigidObjectCollectionDataProperties: + """Comprehensive tests for MockRigidObjectCollectionData properties.""" + + @pytest.fixture + def data(self): + """Create MockRigidObjectCollectionData fixture.""" + return MockRigidObjectCollectionData( + num_instances=4, + num_bodies=5, + device="cpu", + ) + + @pytest.mark.parametrize( + "property_name,expected_shape", + [ + ("body_link_pose_w", (4, 5, 7)), + ("body_link_vel_w", (4, 5, 6)), + ("body_link_state_w", (4, 5, 13)), + ("body_link_pos_w", (4, 5, 3)), + ("body_link_quat_w", (4, 5, 4)), + ("body_link_lin_vel_w", (4, 5, 3)), + ("body_link_ang_vel_w", (4, 5, 3)), + ("body_link_lin_vel_b", (4, 5, 3)), + ("body_link_ang_vel_b", (4, 5, 3)), + ("body_com_pose_w", (4, 5, 7)), + ("body_com_vel_w", (4, 5, 6)), + ("body_com_state_w", (4, 5, 13)), + ("body_com_acc_w", (4, 5, 6)), + ("body_com_pos_w", (4, 5, 3)), + ("body_com_quat_w", (4, 5, 4)), + ("body_com_lin_vel_w", (4, 5, 3)), + ("body_com_ang_vel_w", (4, 5, 3)), + ("body_com_lin_vel_b", (4, 5, 3)), + ("body_com_ang_vel_b", (4, 5, 3)), + ("body_com_lin_acc_w", (4, 5, 3)), + ("body_com_ang_acc_w", (4, 5, 3)), + ("body_com_pose_b", (4, 5, 7)), + ("body_com_pos_b", (4, 5, 3)), + ("body_com_quat_b", (4, 5, 4)), + ("body_mass", (4, 5)), + ("body_inertia", (4, 5, 9)), + ("projected_gravity_b", (4, 5, 3)), + ("heading_w", (4, 5)), + ("default_body_pose", (4, 5, 7)), + ("default_body_vel", (4, 5, 6)), + ("default_body_state", (4, 5, 13)), + ], + ) + def test_property_shapes(self, data, property_name, expected_shape): + """Test that all properties return tensors with correct shapes.""" + prop = getattr(data, property_name) + assert prop.shape == expected_shape, f"{property_name} has wrong shape" + + +# ============================================================================== +# Device Handling Tests +# ============================================================================== + + +class TestDeviceHandling: + """Test that tensors are created on the correct device.""" + + def test_imu_data_device(self): + """Test IMU data tensors are on correct device.""" + data = MockImuData(num_instances=2, device="cpu") + assert data.pos_w.device.type == "cpu" + assert data.quat_w.device.type == "cpu" + + def test_contact_sensor_data_device(self): + """Test contact sensor data tensors are on correct device.""" + data = MockContactSensorData(num_instances=2, num_bodies=3, device="cpu") + assert data.net_forces_w.device.type == "cpu" + + def test_articulation_data_device(self): + """Test articulation data tensors are on correct device.""" + data = MockArticulationData(num_instances=2, num_joints=6, num_bodies=7, device="cpu") + assert data.joint_pos.device.type == "cpu" + assert data.root_link_pose_w.device.type == "cpu" + + def test_setter_moves_tensor_to_device(self): + """Test that setters move tensors to the correct device.""" + data = MockImuData(num_instances=2, device="cpu") + # Create tensor on CPU (default) + test_tensor = torch.randn(2, 3) + data.set_pos_w(test_tensor) + assert data.pos_w.device.type == "cpu" + + +# ============================================================================== +# Composite Property Tests +# ============================================================================== + + +class TestCompositeProperties: + """Test that composite properties are correctly composed from components.""" + + def test_imu_pose_composition(self): + """Test IMU pose_w is correctly composed from pos_w and quat_w.""" + data = MockImuData(num_instances=2, device="cpu") + pos = torch.randn(2, 3) + quat = torch.tensor([[1, 0, 0, 0], [0.707, 0.707, 0, 0]], dtype=torch.float32) + + data.set_pos_w(pos) + data.set_quat_w(quat) + + pose = data.pose_w + assert torch.allclose(pose[:, :3], pos) + assert torch.allclose(pose[:, 3:], quat) + + def test_articulation_root_state_composition(self): + """Test articulation root_link_state_w is correctly composed.""" + data = MockArticulationData(num_instances=2, num_joints=6, num_bodies=7, device="cpu") + pose = torch.randn(2, 7) + vel = torch.randn(2, 6) + + data.set_root_link_pose_w(pose) + data.set_root_link_vel_w(vel) + + state = data.root_link_state_w + assert torch.allclose(state[:, :7], pose) + assert torch.allclose(state[:, 7:], vel) + + def test_articulation_default_state_composition(self): + """Test articulation default_root_state is correctly composed.""" + data = MockArticulationData(num_instances=2, num_joints=6, num_bodies=7, device="cpu") + pose = torch.randn(2, 7) + vel = torch.randn(2, 6) + + data.set_default_root_pose(pose) + data.set_default_root_vel(vel) + + state = data.default_root_state + assert torch.allclose(state[:, :7], pose) + assert torch.allclose(state[:, 7:], vel) + + +# ============================================================================== +# Convenience Alias Property Tests +# ============================================================================== + + +class TestArticulationConvenienceAliases: + """Test convenience alias properties for MockArticulationData.""" + + @pytest.fixture + def data(self): + """Create MockArticulationData fixture.""" + return MockArticulationData( + num_instances=4, + num_joints=12, + num_bodies=13, + num_fixed_tendons=2, + device="cpu", + ) + + # -- Root state aliases (without _link_ or _com_ prefix) -- + @pytest.mark.parametrize( + "alias,source,expected_shape", + [ + ("root_pos_w", "root_link_pos_w", (4, 3)), + ("root_quat_w", "root_link_quat_w", (4, 4)), + ("root_pose_w", "root_link_pose_w", (4, 7)), + ("root_vel_w", "root_com_vel_w", (4, 6)), + ("root_lin_vel_w", "root_com_lin_vel_w", (4, 3)), + ("root_ang_vel_w", "root_com_ang_vel_w", (4, 3)), + ("root_lin_vel_b", "root_com_lin_vel_b", (4, 3)), + ("root_ang_vel_b", "root_com_ang_vel_b", (4, 3)), + ], + ) + def test_root_aliases(self, data, alias, source, expected_shape): + """Test root state aliases reference correct properties.""" + alias_prop = getattr(data, alias) + source_prop = getattr(data, source) + assert alias_prop.shape == expected_shape, f"{alias} has wrong shape" + assert torch.equal(alias_prop, source_prop), f"{alias} should equal {source}" + + # -- Body state aliases (without _link_ or _com_ prefix) -- + @pytest.mark.parametrize( + "alias,source,expected_shape", + [ + ("body_pos_w", "body_link_pos_w", (4, 13, 3)), + ("body_quat_w", "body_link_quat_w", (4, 13, 4)), + ("body_pose_w", "body_link_pose_w", (4, 13, 7)), + ("body_vel_w", "body_com_vel_w", (4, 13, 6)), + ("body_state_w", "body_com_state_w", (4, 13, 13)), + ("body_lin_vel_w", "body_com_lin_vel_w", (4, 13, 3)), + ("body_ang_vel_w", "body_com_ang_vel_w", (4, 13, 3)), + ("body_acc_w", "body_com_acc_w", (4, 13, 6)), + ("body_lin_acc_w", "body_com_lin_acc_w", (4, 13, 3)), + ("body_ang_acc_w", "body_com_ang_acc_w", (4, 13, 3)), + ], + ) + def test_body_aliases(self, data, alias, source, expected_shape): + """Test body state aliases reference correct properties.""" + alias_prop = getattr(data, alias) + source_prop = getattr(data, source) + assert alias_prop.shape == expected_shape, f"{alias} has wrong shape" + assert torch.equal(alias_prop, source_prop), f"{alias} should equal {source}" + + # -- CoM in body frame -- + @pytest.mark.parametrize( + "alias,expected_shape", + [ + ("com_pos_b", (4, 3)), + ("com_quat_b", (4, 4)), + ], + ) + def test_com_body_frame_aliases(self, data, alias, expected_shape): + """Test CoM in body frame aliases.""" + prop = getattr(data, alias) + assert prop.shape == expected_shape, f"{alias} has wrong shape" + + # -- Joint property aliases -- + @pytest.mark.parametrize( + "alias,source,expected_shape", + [ + ("joint_limits", "joint_pos_limits", (4, 12, 2)), + ("joint_velocity_limits", "joint_vel_limits", (4, 12)), + ("joint_friction", "joint_friction_coeff", (4, 12)), + ], + ) + def test_joint_aliases(self, data, alias, source, expected_shape): + """Test joint property aliases.""" + alias_prop = getattr(data, alias) + source_prop = getattr(data, source) + assert alias_prop.shape == expected_shape, f"{alias} has wrong shape" + assert torch.equal(alias_prop, source_prop), f"{alias} should equal {source}" + + # -- Fixed tendon alias -- + def test_fixed_tendon_limit_alias(self, data): + """Test fixed_tendon_limit is alias for fixed_tendon_pos_limits.""" + assert data.fixed_tendon_limit.shape == (4, 2, 2) + assert torch.equal(data.fixed_tendon_limit, data.fixed_tendon_pos_limits) + + +class TestRigidObjectConvenienceAliases: + """Test convenience alias properties for MockRigidObjectData.""" + + @pytest.fixture + def data(self): + """Create MockRigidObjectData fixture.""" + return MockRigidObjectData(num_instances=4, device="cpu") + + @pytest.mark.parametrize( + "alias,source,expected_shape", + [ + ("root_pos_w", "root_link_pos_w", (4, 3)), + ("root_quat_w", "root_link_quat_w", (4, 4)), + ("root_pose_w", "root_link_pose_w", (4, 7)), + ("root_vel_w", "root_com_vel_w", (4, 6)), + ("root_lin_vel_w", "root_com_lin_vel_w", (4, 3)), + ("root_ang_vel_w", "root_com_ang_vel_w", (4, 3)), + ("root_lin_vel_b", "root_com_lin_vel_b", (4, 3)), + ("root_ang_vel_b", "root_com_ang_vel_b", (4, 3)), + ("body_pos_w", "body_link_pos_w", (4, 1, 3)), + ("body_quat_w", "body_link_quat_w", (4, 1, 4)), + ("body_pose_w", "body_link_pose_w", (4, 1, 7)), + ("body_vel_w", "body_com_vel_w", (4, 1, 6)), + ("body_state_w", "body_com_state_w", (4, 1, 13)), + ("body_lin_vel_w", "body_com_lin_vel_w", (4, 1, 3)), + ("body_ang_vel_w", "body_com_ang_vel_w", (4, 1, 3)), + ("body_acc_w", "body_com_acc_w", (4, 1, 6)), + ("body_lin_acc_w", "body_com_lin_acc_w", (4, 1, 3)), + ("body_ang_acc_w", "body_com_ang_acc_w", (4, 1, 3)), + ], + ) + def test_aliases(self, data, alias, source, expected_shape): + """Test convenience aliases reference correct properties.""" + alias_prop = getattr(data, alias) + source_prop = getattr(data, source) + assert alias_prop.shape == expected_shape, f"{alias} has wrong shape" + assert torch.equal(alias_prop, source_prop), f"{alias} should equal {source}" + + @pytest.mark.parametrize( + "alias,expected_shape", + [ + ("com_pos_b", (4, 3)), + ("com_quat_b", (4, 4)), + ], + ) + def test_com_body_frame_aliases(self, data, alias, expected_shape): + """Test CoM in body frame aliases.""" + prop = getattr(data, alias) + assert prop.shape == expected_shape, f"{alias} has wrong shape" + + +class TestRigidObjectCollectionConvenienceAliases: + """Test convenience alias properties for MockRigidObjectCollectionData.""" + + @pytest.fixture + def data(self): + """Create MockRigidObjectCollectionData fixture.""" + return MockRigidObjectCollectionData( + num_instances=4, + num_bodies=5, + device="cpu", + ) + + def test_body_state_w_alias(self, data): + """Test body_state_w is alias for body_com_state_w.""" + assert data.body_state_w.shape == (4, 5, 13) + assert torch.equal(data.body_state_w, data.body_com_state_w) diff --git a/source/isaaclab/test/test_mock_interfaces/test_mock_sensors.py b/source/isaaclab/test/test_mock_interfaces/test_mock_sensors.py new file mode 100644 index 00000000000..42dbfb3c635 --- /dev/null +++ b/source/isaaclab/test/test_mock_interfaces/test_mock_sensors.py @@ -0,0 +1,309 @@ +# 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 mock sensor interfaces.""" + +import pytest +import torch + +from isaaclab.test.mock_interfaces.sensors import ( + MockContactSensor, + MockFrameTransformer, + MockImu, + create_mock_foot_contact_sensor, + create_mock_frame_transformer, + create_mock_imu, +) +from isaaclab.test.mock_interfaces.utils import MockSensorBuilder + +# ============================================================================== +# MockImu Tests +# ============================================================================== + + +class TestMockImu: + """Tests for MockImu and MockImuData.""" + + @pytest.fixture + def imu(self): + """Create a mock IMU fixture.""" + return MockImu(num_instances=4, device="cpu") + + def test_initialization(self, imu): + """Test that MockImu initializes correctly.""" + assert imu.num_instances == 4 + assert imu.device == "cpu" + assert imu.data is not None + + def test_lazy_tensor_initialization(self, imu): + """Test that unset properties return zero tensors with correct shapes.""" + # Position + pos = imu.data.pos_w + assert pos.shape == (4, 3) + assert torch.all(pos == 0) + + # Quaternion (should be identity: w=1, x=0, y=0, z=0) + quat = imu.data.quat_w + assert quat.shape == (4, 4) + assert torch.all(quat[:, 0] == 1) # w component + assert torch.all(quat[:, 1:] == 0) # xyz components + + # Velocities and accelerations + assert imu.data.lin_vel_b.shape == (4, 3) + assert imu.data.ang_vel_b.shape == (4, 3) + assert imu.data.lin_acc_b.shape == (4, 3) + assert imu.data.ang_acc_b.shape == (4, 3) + + def test_projected_gravity_default(self, imu): + """Test default gravity direction.""" + gravity = imu.data.projected_gravity_b + assert gravity.shape == (4, 3) + # Default gravity should point down: (0, 0, -1) + assert torch.all(gravity[:, 2] == -1) + + def test_set_mock_data(self, imu): + """Test bulk data setter.""" + lin_vel = torch.randn(4, 3) + ang_vel = torch.randn(4, 3) + + imu.data.set_mock_data(lin_vel_b=lin_vel, ang_vel_b=ang_vel) + + assert torch.allclose(imu.data.lin_vel_b, lin_vel) + assert torch.allclose(imu.data.ang_vel_b, ang_vel) + + def test_per_property_setter(self, imu): + """Test individual property setters.""" + lin_acc = torch.randn(4, 3) + imu.data.set_lin_acc_b(lin_acc) + assert torch.allclose(imu.data.lin_acc_b, lin_acc) + + def test_pose_composition(self, imu): + """Test that pose_w combines pos_w and quat_w correctly.""" + pos = torch.randn(4, 3) + quat = torch.tensor([[1, 0, 0, 0]] * 4, dtype=torch.float32) + + imu.data.set_pos_w(pos) + imu.data.set_quat_w(quat) + + pose = imu.data.pose_w + assert pose.shape == (4, 7) + assert torch.allclose(pose[:, :3], pos) + assert torch.allclose(pose[:, 3:], quat) + + +# ============================================================================== +# MockContactSensor Tests +# ============================================================================== + + +class TestMockContactSensor: + """Tests for MockContactSensor and MockContactSensorData.""" + + @pytest.fixture + def sensor(self): + """Create a mock contact sensor fixture.""" + return MockContactSensor( + num_instances=4, + num_bodies=4, + body_names=["FL_foot", "FR_foot", "RL_foot", "RR_foot"], + device="cpu", + ) + + def test_initialization(self, sensor): + """Test that MockContactSensor initializes correctly.""" + assert sensor.num_instances == 4 + assert sensor.num_bodies == 4 + assert sensor.body_names == ["FL_foot", "FR_foot", "RL_foot", "RR_foot"] + assert sensor.contact_view is None + + def test_lazy_tensor_shapes(self, sensor): + """Test that unset properties return tensors with correct shapes.""" + forces = sensor.data.net_forces_w + assert forces.shape == (4, 4, 3) + + contact_time = sensor.data.current_contact_time + assert contact_time.shape == (4, 4) + + air_time = sensor.data.current_air_time + assert air_time.shape == (4, 4) + + def test_find_bodies(self, sensor): + """Test body finding by regex.""" + # Find all bodies matching ".*_foot" + indices, names = sensor.find_bodies(".*_foot") + assert len(indices) == 4 + assert set(names) == {"FL_foot", "FR_foot", "RL_foot", "RR_foot"} + + # Find specific body + indices, names = sensor.find_bodies("FL_foot") + assert indices == [0] + assert names == ["FL_foot"] + + # Find front feet + indices, names = sensor.find_bodies("F._foot") + assert len(indices) == 2 + assert "FL_foot" in names + assert "FR_foot" in names + + def test_compute_first_contact(self, sensor): + """Test first contact computation.""" + # Set contact time to 0.5 for all bodies + sensor.data.set_current_contact_time(torch.full((4, 4), 0.5)) + + # Check with dt=1.0 - should be True (0.5 < 1.0) + first_contact = sensor.compute_first_contact(dt=1.0) + assert torch.all(first_contact) + + # Check with dt=0.1 - should be False (0.5 > 0.1) + first_contact = sensor.compute_first_contact(dt=0.1) + assert torch.all(~first_contact) + + def test_compute_first_air(self, sensor): + """Test first air computation.""" + sensor.data.set_current_air_time(torch.full((4, 4), 0.2)) + + first_air = sensor.compute_first_air(dt=0.5) + assert torch.all(first_air) + + def test_history_buffer(self): + """Test history buffer when enabled.""" + sensor_with_history = MockContactSensor( + num_instances=2, + num_bodies=2, + history_length=3, + device="cpu", + ) + + history = sensor_with_history.data.net_forces_w_history + assert history is not None + assert history.shape == (2, 3, 2, 3) + + def test_no_history_buffer(self, sensor): + """Test that history buffer is None when not enabled.""" + history = sensor.data.net_forces_w_history + assert history is None + + +# ============================================================================== +# MockFrameTransformer Tests +# ============================================================================== + + +class TestMockFrameTransformer: + """Tests for MockFrameTransformer and MockFrameTransformerData.""" + + @pytest.fixture + def transformer(self): + """Create a mock frame transformer fixture.""" + return MockFrameTransformer( + num_instances=2, + num_target_frames=3, + target_frame_names=["target_1", "target_2", "target_3"], + device="cpu", + ) + + def test_initialization(self, transformer): + """Test that MockFrameTransformer initializes correctly.""" + assert transformer.num_instances == 2 + assert transformer.num_bodies == 3 + assert transformer.body_names == ["target_1", "target_2", "target_3"] + + def test_data_shapes(self, transformer): + """Test that data properties have correct shapes.""" + # Source frame + assert transformer.data.source_pos_w.shape == (2, 3) + assert transformer.data.source_quat_w.shape == (2, 4) + + # Target frames + assert transformer.data.target_pos_w.shape == (2, 3, 3) + assert transformer.data.target_quat_w.shape == (2, 3, 4) + + # Relative poses + assert transformer.data.target_pos_source.shape == (2, 3, 3) + assert transformer.data.target_quat_source.shape == (2, 3, 4) + + def test_pose_composition(self, transformer): + """Test that pose properties combine position and orientation correctly.""" + source_pose = transformer.data.source_pose_w + assert source_pose.shape == (2, 7) + + target_pose = transformer.data.target_pose_w + assert target_pose.shape == (2, 3, 7) + + def test_find_bodies(self, transformer): + """Test frame finding by regex.""" + indices, names = transformer.find_bodies("target_.*") + assert len(indices) == 3 + + indices, names = transformer.find_bodies("target_1") + assert indices == [0] + + +# ============================================================================== +# Factory Function Tests +# ============================================================================== + + +class TestSensorFactories: + """Tests for sensor factory functions.""" + + def test_create_mock_imu(self): + """Test IMU factory function.""" + imu = create_mock_imu(num_instances=4) + assert imu.num_instances == 4 + assert imu.data.projected_gravity_b.shape == (4, 3) + + def test_create_mock_foot_contact_sensor(self): + """Test foot contact sensor factory function.""" + sensor = create_mock_foot_contact_sensor(num_instances=4, num_feet=4) + assert sensor.num_instances == 4 + assert sensor.num_bodies == 4 + assert "FL_foot" in sensor.body_names + + def test_create_mock_frame_transformer(self): + """Test frame transformer factory function.""" + transformer = create_mock_frame_transformer(num_instances=2, num_target_frames=5) + assert transformer.num_instances == 2 + assert transformer.num_bodies == 5 + + +# ============================================================================== +# MockSensorBuilder Tests +# ============================================================================== + + +class TestMockSensorBuilder: + """Tests for MockSensorBuilder.""" + + def test_build_contact_sensor(self): + """Test building a contact sensor.""" + sensor = ( + MockSensorBuilder("contact") + .with_num_instances(4) + .with_bodies(["foot_1", "foot_2"]) + .with_history_length(3) + .build() + ) + + assert sensor.num_instances == 4 + assert sensor.num_bodies == 2 + assert sensor.data.net_forces_w_history.shape == (4, 3, 2, 3) + + def test_build_imu_sensor(self): + """Test building an IMU sensor.""" + sensor = MockSensorBuilder("imu").with_num_instances(2).build() + assert sensor.num_instances == 2 + + def test_build_frame_transformer(self): + """Test building a frame transformer sensor.""" + sensor = ( + MockSensorBuilder("frame_transformer") + .with_num_instances(3) + .with_target_frames(["target_1", "target_2"]) + .build() + ) + + assert sensor.num_instances == 3 + assert sensor.num_bodies == 2 diff --git a/source/isaaclab_contrib/config/extension.toml b/source/isaaclab_contrib/config/extension.toml index e31caf78281..15b76c8b510 100644 --- a/source/isaaclab_contrib/config/extension.toml +++ b/source/isaaclab_contrib/config/extension.toml @@ -1,6 +1,6 @@ [package] # Semantic Versioning is used: https://semver.org/ -version = "0.2.0" +version = "0.2.1" # Description title = "Isaac Lab External Contributions" diff --git a/source/isaaclab_contrib/docs/CHANGELOG.rst b/source/isaaclab_contrib/docs/CHANGELOG.rst index 20859db98cd..babd2dcd734 100644 --- a/source/isaaclab_contrib/docs/CHANGELOG.rst +++ b/source/isaaclab_contrib/docs/CHANGELOG.rst @@ -1,6 +1,15 @@ Changelog --------- +0.2.1 (2026-02-03) +~~~~~~~~~~~~~~~~~~ + +Changed +^^^^^^^ + +* Updated the multirotor asset to use the new base classes from the isaaclab_physx package. + + 0.2.0 (2026-01-30) ~~~~~~~~~~~~~~~~~~ diff --git a/source/isaaclab_physx/benchmark/assets/benchmark_articulation.py b/source/isaaclab_physx/benchmark/assets/benchmark_articulation.py new file mode 100644 index 00000000000..3beef7ea158 --- /dev/null +++ b/source/isaaclab_physx/benchmark/assets/benchmark_articulation.py @@ -0,0 +1,1136 @@ +# 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 + +"""Micro-benchmarking framework for Articulation class (PhysX backend). + +This module provides a benchmarking framework to measure the performance of setter and writer +methods in the Articulation class. Each method is benchmarked under two scenarios: + +1. **Torch List**: Inputs are PyTorch tensors with list indices. +2. **Torch Tensor**: Inputs are PyTorch tensors with tensor indices. + +Usage: + python benchmark_articulation.py [--num_iterations N] [--warmup_steps W] + [--num_instances I] [--num_bodies B] [--num_joints J] + +Example: + python benchmark_articulation.py --num_iterations 1000 --warmup_steps 10 + python benchmark_articulation.py --mode torch_list # Only run list-based benchmarks + python benchmark_articulation.py --mode torch_tensor # Only run tensor-based benchmarks +""" + +from __future__ import annotations + +import argparse +import sys +import warnings +from types import ModuleType +from unittest.mock import MagicMock + +import torch +import warp as wp + +# Initialize Warp +wp.init() + + +# ============================================================================= +# Mock Setup - Must happen BEFORE importing Articulation +# ============================================================================= + + +class MockPhysicsSimView: + """Simple mock for the physics simulation view.""" + + def get_gravity(self): + return (0.0, 0.0, -9.81) + + def update_articulations_kinematic(self): + pass + + +class MockSimulationManager: + """Simple mock for SimulationManager.""" + + @staticmethod + def get_physics_sim_view(): + return MockPhysicsSimView() + + +# Mock isaacsim.core.simulation_manager +mock_sim_manager_module = ModuleType("isaacsim.core.simulation_manager") +mock_sim_manager_module.SimulationManager = MockSimulationManager +sys.modules["isaacsim"] = ModuleType("isaacsim") +sys.modules["isaacsim.core"] = ModuleType("isaacsim.core") +sys.modules["isaacsim.core.simulation_manager"] = mock_sim_manager_module + +# Mock pxr (USD library) +sys.modules["pxr"] = MagicMock() +sys.modules["pxr.Usd"] = MagicMock() +sys.modules["pxr.UsdGeom"] = MagicMock() +sys.modules["pxr.UsdPhysics"] = MagicMock() +sys.modules["pxr.PhysxSchema"] = MagicMock() + +# Mock omni module hierarchy (must be ModuleType for proper package behavior) +omni_mocks = [ + "omni", + "omni.kit", + "omni.kit.app", + "omni.kit.commands", + "omni.usd", + "omni.client", + "omni.timeline", + "omni.physx", + "omni.physx.scripts", + "omni.physx.scripts.utils", + "omni.physics", + "omni.physics.tensors", + "omni.physics.tensors.impl", + "omni.physics.tensors.impl.api", +] +for mod_name in omni_mocks: + mock = MagicMock() + mock.__name__ = mod_name + mock.__path__ = [] + mock.__package__ = mod_name + sys.modules[mod_name] = mock + +# Mock carb (needed by isaaclab.utils.assets) +mock_carb = MagicMock() +mock_carb.settings.get_settings.return_value.get.return_value = "/mock/path" +sys.modules["carb"] = mock_carb + +# Mock isaaclab.sim module hierarchy (to avoid importing converters) +sim_mock = MagicMock() +sim_mock.find_first_matching_prim = MagicMock() +sim_mock.get_all_matching_child_prims = MagicMock(return_value=[]) +sys.modules["isaaclab.sim"] = sim_mock +sys.modules["isaaclab.sim.utils"] = MagicMock() +sys.modules["isaaclab.sim.utils.stage"] = MagicMock() +sys.modules["isaaclab.sim.utils.queries"] = MagicMock() +sys.modules["isaaclab.sim.converters"] = MagicMock() + +# Mock prettytable +sys.modules["prettytable"] = MagicMock() + +# Mock WrenchComposer - import from mock_interfaces and patch into the module +from isaaclab.test.mock_interfaces.utils import MockWrenchComposer + +mock_wrench_module = ModuleType("isaaclab.utils.wrench_composer") +mock_wrench_module.WrenchComposer = MockWrenchComposer +sys.modules["isaaclab.utils.wrench_composer"] = mock_wrench_module + + +# Mock base classes to avoid importing full isaaclab.assets package +class BaseArticulation: + """Mock base class.""" + + @property + def device(self) -> str: + return self._device + + +class BaseArticulationData: + """Mock base class.""" + + def __init__(self, root_view, device: str): + self.device = device + + +mock_base_articulation = ModuleType("isaaclab.assets.articulation.base_articulation") +mock_base_articulation.BaseArticulation = BaseArticulation +sys.modules["isaaclab.assets.articulation.base_articulation"] = mock_base_articulation + +mock_base_articulation_data = ModuleType("isaaclab.assets.articulation.base_articulation_data") +mock_base_articulation_data.BaseArticulationData = BaseArticulationData +sys.modules["isaaclab.assets.articulation.base_articulation_data"] = mock_base_articulation_data + +# Mock ArticulationCfg +mock_cfg_module = ModuleType("isaaclab.assets.articulation.articulation_cfg") + + +class ArticulationCfg: + """Mock ArticulationCfg for testing.""" + + def __init__( + self, prim_path: str = "/World/Robot", soft_joint_pos_limit_factor: float = 1.0, actuators: dict = None + ): + self.prim_path = prim_path + self.soft_joint_pos_limit_factor = soft_joint_pos_limit_factor + self.actuators = actuators or {} + + +mock_cfg_module.ArticulationCfg = ArticulationCfg +sys.modules["isaaclab.assets.articulation.articulation_cfg"] = mock_cfg_module + +# Mock actuators module +mock_actuators = ModuleType("isaaclab.actuators") +mock_actuators.ActuatorBase = MagicMock() +mock_actuators.ActuatorBaseCfg = MagicMock() +mock_actuators.ImplicitActuator = MagicMock() +sys.modules["isaaclab.actuators"] = mock_actuators + +# Mock utils modules - need to include all string functions used by isaaclab.utils.dict +mock_string_utils = ModuleType("isaaclab.utils.string") +mock_string_utils.resolve_matching_names = MagicMock(return_value=([], [])) +mock_string_utils.callable_to_string = MagicMock(return_value="") +mock_string_utils.string_to_callable = MagicMock(return_value=None) +mock_string_utils.string_to_slice = MagicMock(return_value=slice(None)) +sys.modules["isaaclab.utils.string"] = mock_string_utils + +mock_types = ModuleType("isaaclab.utils.types") +mock_types.ArticulationActions = MagicMock() +sys.modules["isaaclab.utils.types"] = mock_types + +mock_version = ModuleType("isaaclab.utils.version") + + +# Create a proper version tuple class that has .major, .minor, .patch attributes +class VersionTuple: + def __init__(self, major, minor, patch): + self.major = major + self.minor = minor + self.patch = patch + + def __iter__(self): + return iter((self.major, self.minor, self.patch)) + + +mock_version.get_isaac_sim_version = MagicMock(return_value=VersionTuple(4, 5, 0)) +sys.modules["isaaclab.utils.version"] = mock_version + +# Now import via importlib to bypass __init__.py +import importlib.util +from pathlib import Path + +benchmark_dir = Path(__file__).resolve().parent + +# Load ArticulationData +articulation_data_path = ( + benchmark_dir.parents[1] / "isaaclab_physx" / "assets" / "articulation" / "articulation_data.py" +) +spec = importlib.util.spec_from_file_location( + "isaaclab_physx.assets.articulation.articulation_data", articulation_data_path +) +articulation_data_module = importlib.util.module_from_spec(spec) +sys.modules["isaaclab_physx.assets.articulation.articulation_data"] = articulation_data_module +spec.loader.exec_module(articulation_data_module) +ArticulationData = articulation_data_module.ArticulationData + +# Load Articulation +articulation_path = benchmark_dir.parents[1] / "isaaclab_physx" / "assets" / "articulation" / "articulation.py" +spec = importlib.util.spec_from_file_location("isaaclab_physx.assets.articulation.articulation", articulation_path) +articulation_module = importlib.util.module_from_spec(spec) +sys.modules["isaaclab_physx.assets.articulation.articulation"] = articulation_module +spec.loader.exec_module(articulation_module) +Articulation = articulation_module.Articulation + +# Import shared utilities from common module +# Import mock classes from PhysX test utilities +from isaaclab_physx.test.mock_interfaces.views import MockArticulationView + +from isaaclab.test.benchmark import ( + BenchmarkConfig, + MethodBenchmark, + benchmark_method, + export_results_csv, + export_results_json, + get_default_output_filename, + get_hardware_info, + make_tensor_body_ids, + make_tensor_env_ids, + make_tensor_joint_ids, + print_hardware_info, + print_results, +) + +# Suppress deprecation warnings during benchmarking +warnings.filterwarnings("ignore", category=DeprecationWarning) +warnings.filterwarnings("ignore", category=UserWarning) + +# Also suppress logging warnings (the body_acc_w deprecation warnings use logging) +import logging + +logging.getLogger("isaaclab_physx").setLevel(logging.ERROR) +logging.getLogger("isaaclab").setLevel(logging.ERROR) + + +def create_test_articulation( + num_instances: int = 2, + num_joints: int = 6, + num_bodies: int = 7, + device: str = "cuda:0", +) -> tuple[Articulation, MockArticulationView, MagicMock]: + """Create a test Articulation instance with mocked dependencies.""" + joint_names = [f"joint_{i}" for i in range(num_joints)] + body_names = [f"body_{i}" for i in range(num_bodies)] + + articulation = object.__new__(Articulation) + + articulation.cfg = ArticulationCfg( + prim_path="/World/Robot", + soft_joint_pos_limit_factor=1.0, + actuators={}, + ) + + # Create PhysX mock view + mock_view = MockArticulationView( + count=num_instances, + num_links=num_bodies, + num_dofs=num_joints, + device=device, + ) + mock_view.set_random_mock_data() + + # Set up the mock view's metatype for accessing names/counts + mock_metatype = MagicMock() + mock_metatype.fixed_base = False + mock_metatype.dof_count = num_joints + mock_metatype.link_count = num_bodies + mock_metatype.dof_names = joint_names + mock_metatype.link_names = body_names + object.__setattr__(mock_view, "_shared_metatype", mock_metatype) + + object.__setattr__(articulation, "_root_view", mock_view) + object.__setattr__(articulation, "_device", device) + + # Create ArticulationData instance (SimulationManager already mocked at module level) + data = ArticulationData(mock_view, device) + object.__setattr__(articulation, "_data", data) + + # Create mock wrench composers (pass articulation which has num_instances, num_bodies, device properties) + mock_inst_wrench = MockWrenchComposer(articulation) + mock_perm_wrench = MockWrenchComposer(articulation) + object.__setattr__(articulation, "_instantaneous_wrench_composer", mock_inst_wrench) + object.__setattr__(articulation, "_permanent_wrench_composer", mock_perm_wrench) + + # Set up other required attributes + object.__setattr__(articulation, "actuators", {}) + object.__setattr__(articulation, "_has_implicit_actuators", False) + object.__setattr__(articulation, "_ALL_INDICES", torch.arange(num_instances, device=device)) + object.__setattr__(articulation, "_ALL_BODY_INDICES", torch.arange(num_bodies, device=device)) + object.__setattr__(articulation, "_ALL_JOINT_INDICES", torch.arange(num_joints, device=device)) + + # Warp arrays for set_external_force_and_torque + all_indices = torch.arange(num_instances, dtype=torch.int32, device=device) + all_body_indices = torch.arange(num_bodies, dtype=torch.int32, device=device) + object.__setattr__(articulation, "_ALL_INDICES_WP", wp.from_torch(all_indices, dtype=wp.int32)) + object.__setattr__(articulation, "_ALL_BODY_INDICES_WP", wp.from_torch(all_body_indices, dtype=wp.int32)) + + # Initialize joint targets + object.__setattr__(articulation, "_joint_pos_target_sim", torch.zeros(num_instances, num_joints, device=device)) + object.__setattr__(articulation, "_joint_vel_target_sim", torch.zeros(num_instances, num_joints, device=device)) + object.__setattr__(articulation, "_joint_effort_target_sim", torch.zeros(num_instances, num_joints, device=device)) + + return articulation, mock_view, None + + +# ============================================================================= +# Input Generators (Torch-only for PhysX backend) +# ============================================================================= + + +# --- Root Link Pose --- +def gen_root_link_pose_torch_list(config: BenchmarkConfig) -> dict: + """Generate Torch inputs with list env_ids for write_root_link_pose_to_sim.""" + return { + "root_pose": torch.rand(config.num_instances, 7, device=config.device, dtype=torch.float32), + "env_ids": list(range(config.num_instances)), + } + + +def gen_root_link_pose_torch_tensor(config: BenchmarkConfig) -> dict: + """Generate Torch inputs with tensor env_ids for write_root_link_pose_to_sim.""" + return { + "root_pose": torch.rand(config.num_instances, 7, device=config.device, dtype=torch.float32), + "env_ids": make_tensor_env_ids(config.num_instances, config.device), + } + + +# --- Root COM Pose --- +def gen_root_com_pose_torch_list(config: BenchmarkConfig) -> dict: + """Generate Torch inputs with list env_ids for write_root_com_pose_to_sim.""" + return { + "root_pose": torch.rand(config.num_instances, 7, device=config.device, dtype=torch.float32), + "env_ids": list(range(config.num_instances)), + } + + +def gen_root_com_pose_torch_tensor(config: BenchmarkConfig) -> dict: + """Generate Torch inputs with tensor env_ids for write_root_com_pose_to_sim.""" + return { + "root_pose": torch.rand(config.num_instances, 7, device=config.device, dtype=torch.float32), + "env_ids": make_tensor_env_ids(config.num_instances, config.device), + } + + +# --- Root Link Velocity --- +def gen_root_link_velocity_torch_list(config: BenchmarkConfig) -> dict: + """Generate Torch inputs with list env_ids for write_root_link_velocity_to_sim.""" + return { + "root_velocity": torch.rand(config.num_instances, 6, device=config.device, dtype=torch.float32), + "env_ids": list(range(config.num_instances)), + } + + +def gen_root_link_velocity_torch_tensor(config: BenchmarkConfig) -> dict: + """Generate Torch inputs with tensor env_ids for write_root_link_velocity_to_sim.""" + return { + "root_velocity": torch.rand(config.num_instances, 6, device=config.device, dtype=torch.float32), + "env_ids": make_tensor_env_ids(config.num_instances, config.device), + } + + +# --- Root COM Velocity --- +def gen_root_com_velocity_torch_list(config: BenchmarkConfig) -> dict: + """Generate Torch inputs with list env_ids for write_root_com_velocity_to_sim.""" + return { + "root_velocity": torch.rand(config.num_instances, 6, device=config.device, dtype=torch.float32), + "env_ids": list(range(config.num_instances)), + } + + +def gen_root_com_velocity_torch_tensor(config: BenchmarkConfig) -> dict: + """Generate Torch inputs with tensor env_ids for write_root_com_velocity_to_sim.""" + return { + "root_velocity": torch.rand(config.num_instances, 6, device=config.device, dtype=torch.float32), + "env_ids": make_tensor_env_ids(config.num_instances, config.device), + } + + +# --- Root State (Deprecated) --- +def gen_root_state_torch_list(config: BenchmarkConfig) -> dict: + """Generate Torch inputs with list env_ids for write_root_state_to_sim.""" + return { + "root_state": torch.rand(config.num_instances, 13, device=config.device, dtype=torch.float32), + "env_ids": list(range(config.num_instances)), + } + + +def gen_root_state_torch_tensor(config: BenchmarkConfig) -> dict: + """Generate Torch inputs with tensor env_ids for write_root_state_to_sim.""" + return { + "root_state": torch.rand(config.num_instances, 13, device=config.device, dtype=torch.float32), + "env_ids": make_tensor_env_ids(config.num_instances, config.device), + } + + +# --- Root COM State (Deprecated) --- +def gen_root_com_state_torch_list(config: BenchmarkConfig) -> dict: + """Generate Torch inputs with list env_ids for write_root_com_state_to_sim.""" + return { + "root_state": torch.rand(config.num_instances, 13, device=config.device, dtype=torch.float32), + "env_ids": list(range(config.num_instances)), + } + + +def gen_root_com_state_torch_tensor(config: BenchmarkConfig) -> dict: + """Generate Torch inputs with tensor env_ids for write_root_com_state_to_sim.""" + return { + "root_state": torch.rand(config.num_instances, 13, device=config.device, dtype=torch.float32), + "env_ids": make_tensor_env_ids(config.num_instances, config.device), + } + + +# --- Root Link State (Deprecated) --- +def gen_root_link_state_torch_list(config: BenchmarkConfig) -> dict: + """Generate Torch inputs with list env_ids for write_root_link_state_to_sim.""" + return { + "root_state": torch.rand(config.num_instances, 13, device=config.device, dtype=torch.float32), + "env_ids": list(range(config.num_instances)), + } + + +def gen_root_link_state_torch_tensor(config: BenchmarkConfig) -> dict: + """Generate Torch inputs with tensor env_ids for write_root_link_state_to_sim.""" + return { + "root_state": torch.rand(config.num_instances, 13, device=config.device, dtype=torch.float32), + "env_ids": make_tensor_env_ids(config.num_instances, config.device), + } + + +# --- Joint State --- +def gen_joint_state_torch_list(config: BenchmarkConfig) -> dict: + """Generate Torch inputs with list ids for write_joint_state_to_sim.""" + return { + "position": torch.rand(config.num_instances, config.num_joints, device=config.device, dtype=torch.float32), + "velocity": torch.rand(config.num_instances, config.num_joints, device=config.device, dtype=torch.float32), + "env_ids": list(range(config.num_instances)), + "joint_ids": list(range(config.num_joints)), + } + + +def gen_joint_state_torch_tensor(config: BenchmarkConfig) -> dict: + """Generate Torch inputs with tensor ids for write_joint_state_to_sim.""" + return { + "position": torch.rand(config.num_instances, config.num_joints, device=config.device, dtype=torch.float32), + "velocity": torch.rand(config.num_instances, config.num_joints, device=config.device, dtype=torch.float32), + "env_ids": make_tensor_env_ids(config.num_instances, config.device), + "joint_ids": make_tensor_joint_ids(config.num_joints, config.device), + } + + +# --- Joint Position --- +def gen_joint_position_torch_list(config: BenchmarkConfig) -> dict: + """Generate Torch inputs with list ids for write_joint_position_to_sim.""" + return { + "position": torch.rand(config.num_instances, config.num_joints, device=config.device, dtype=torch.float32), + "env_ids": list(range(config.num_instances)), + "joint_ids": list(range(config.num_joints)), + } + + +def gen_joint_position_torch_tensor(config: BenchmarkConfig) -> dict: + """Generate Torch inputs with tensor ids for write_joint_position_to_sim.""" + return { + "position": torch.rand(config.num_instances, config.num_joints, device=config.device, dtype=torch.float32), + "env_ids": make_tensor_env_ids(config.num_instances, config.device), + "joint_ids": make_tensor_joint_ids(config.num_joints, config.device), + } + + +# --- Joint Velocity --- +def gen_joint_velocity_torch_list(config: BenchmarkConfig) -> dict: + """Generate Torch inputs with list ids for write_joint_velocity_to_sim.""" + return { + "velocity": torch.rand(config.num_instances, config.num_joints, device=config.device, dtype=torch.float32), + "env_ids": list(range(config.num_instances)), + "joint_ids": list(range(config.num_joints)), + } + + +def gen_joint_velocity_torch_tensor(config: BenchmarkConfig) -> dict: + """Generate Torch inputs with tensor ids for write_joint_velocity_to_sim.""" + return { + "velocity": torch.rand(config.num_instances, config.num_joints, device=config.device, dtype=torch.float32), + "env_ids": make_tensor_env_ids(config.num_instances, config.device), + "joint_ids": make_tensor_joint_ids(config.num_joints, config.device), + } + + +# --- Joint Stiffness --- +def gen_joint_stiffness_torch_list(config: BenchmarkConfig) -> dict: + """Generate Torch inputs with list ids for write_joint_stiffness_to_sim.""" + return { + "stiffness": torch.rand(config.num_instances, config.num_joints, device=config.device, dtype=torch.float32), + "env_ids": list(range(config.num_instances)), + "joint_ids": list(range(config.num_joints)), + } + + +def gen_joint_stiffness_torch_tensor(config: BenchmarkConfig) -> dict: + """Generate Torch inputs with tensor ids for write_joint_stiffness_to_sim.""" + return { + "stiffness": torch.rand(config.num_instances, config.num_joints, device=config.device, dtype=torch.float32), + "env_ids": make_tensor_env_ids(config.num_instances, config.device), + "joint_ids": make_tensor_joint_ids(config.num_joints, config.device), + } + + +# --- Joint Damping --- +def gen_joint_damping_torch_list(config: BenchmarkConfig) -> dict: + """Generate Torch inputs with list ids for write_joint_damping_to_sim.""" + return { + "damping": torch.rand(config.num_instances, config.num_joints, device=config.device, dtype=torch.float32), + "env_ids": list(range(config.num_instances)), + "joint_ids": list(range(config.num_joints)), + } + + +def gen_joint_damping_torch_tensor(config: BenchmarkConfig) -> dict: + """Generate Torch inputs with tensor ids for write_joint_damping_to_sim.""" + return { + "damping": torch.rand(config.num_instances, config.num_joints, device=config.device, dtype=torch.float32), + "env_ids": make_tensor_env_ids(config.num_instances, config.device), + "joint_ids": make_tensor_joint_ids(config.num_joints, config.device), + } + + +# --- Joint Position Limit --- +def gen_joint_position_limit_torch_list(config: BenchmarkConfig) -> dict: + """Generate Torch inputs with list ids for write_joint_position_limit_to_sim.""" + # limits shape is (N, J, 2) where [:,: ,0] is lower and [:,:,1] is upper + lower = torch.rand(config.num_instances, config.num_joints, 1, device=config.device, dtype=torch.float32) * -3.14 + upper = torch.rand(config.num_instances, config.num_joints, 1, device=config.device, dtype=torch.float32) * 3.14 + return { + "limits": torch.cat([lower, upper], dim=-1), + "env_ids": list(range(config.num_instances)), + "joint_ids": list(range(config.num_joints)), + } + + +def gen_joint_position_limit_torch_tensor(config: BenchmarkConfig) -> dict: + """Generate Torch inputs with tensor ids for write_joint_position_limit_to_sim.""" + # limits shape is (N, J, 2) where [:,: ,0] is lower and [:,:,1] is upper + lower = torch.rand(config.num_instances, config.num_joints, 1, device=config.device, dtype=torch.float32) * -3.14 + upper = torch.rand(config.num_instances, config.num_joints, 1, device=config.device, dtype=torch.float32) * 3.14 + return { + "limits": torch.cat([lower, upper], dim=-1), + "env_ids": make_tensor_env_ids(config.num_instances, config.device), + "joint_ids": make_tensor_joint_ids(config.num_joints, config.device), + } + + +# --- Joint Velocity Limit --- +def gen_joint_velocity_limit_torch_list(config: BenchmarkConfig) -> dict: + """Generate Torch inputs with list ids for write_joint_velocity_limit_to_sim.""" + return { + "limits": torch.rand(config.num_instances, config.num_joints, device=config.device, dtype=torch.float32) * 10.0, + "env_ids": list(range(config.num_instances)), + "joint_ids": list(range(config.num_joints)), + } + + +def gen_joint_velocity_limit_torch_tensor(config: BenchmarkConfig) -> dict: + """Generate Torch inputs with tensor ids for write_joint_velocity_limit_to_sim.""" + return { + "limits": torch.rand(config.num_instances, config.num_joints, device=config.device, dtype=torch.float32) * 10.0, + "env_ids": make_tensor_env_ids(config.num_instances, config.device), + "joint_ids": make_tensor_joint_ids(config.num_joints, config.device), + } + + +# --- Joint Effort Limit --- +def gen_joint_effort_limit_torch_list(config: BenchmarkConfig) -> dict: + """Generate Torch inputs with list ids for write_joint_effort_limit_to_sim.""" + return { + "limits": ( + torch.rand(config.num_instances, config.num_joints, device=config.device, dtype=torch.float32) * 100.0 + ), + "env_ids": list(range(config.num_instances)), + "joint_ids": list(range(config.num_joints)), + } + + +def gen_joint_effort_limit_torch_tensor(config: BenchmarkConfig) -> dict: + """Generate Torch inputs with tensor ids for write_joint_effort_limit_to_sim.""" + return { + "limits": ( + torch.rand(config.num_instances, config.num_joints, device=config.device, dtype=torch.float32) * 100.0 + ), + "env_ids": make_tensor_env_ids(config.num_instances, config.device), + "joint_ids": make_tensor_joint_ids(config.num_joints, config.device), + } + + +# --- Joint Armature --- +def gen_joint_armature_torch_list(config: BenchmarkConfig) -> dict: + """Generate Torch inputs with list ids for write_joint_armature_to_sim.""" + return { + "armature": ( + torch.rand(config.num_instances, config.num_joints, device=config.device, dtype=torch.float32) * 0.1 + ), + "env_ids": list(range(config.num_instances)), + "joint_ids": list(range(config.num_joints)), + } + + +def gen_joint_armature_torch_tensor(config: BenchmarkConfig) -> dict: + """Generate Torch inputs with tensor ids for write_joint_armature_to_sim.""" + return { + "armature": ( + torch.rand(config.num_instances, config.num_joints, device=config.device, dtype=torch.float32) * 0.1 + ), + "env_ids": make_tensor_env_ids(config.num_instances, config.device), + "joint_ids": make_tensor_joint_ids(config.num_joints, config.device), + } + + +# --- Joint Friction Coefficient --- +def gen_joint_friction_coefficient_torch_list(config: BenchmarkConfig) -> dict: + """Generate Torch inputs with list ids for write_joint_friction_coefficient_to_sim.""" + return { + "joint_friction_coeff": ( + torch.rand(config.num_instances, config.num_joints, device=config.device, dtype=torch.float32) * 0.5 + ), + "env_ids": list(range(config.num_instances)), + "joint_ids": list(range(config.num_joints)), + } + + +def gen_joint_friction_coefficient_torch_tensor(config: BenchmarkConfig) -> dict: + """Generate Torch inputs with tensor ids for write_joint_friction_coefficient_to_sim.""" + return { + "joint_friction_coeff": ( + torch.rand(config.num_instances, config.num_joints, device=config.device, dtype=torch.float32) * 0.5 + ), + "env_ids": make_tensor_env_ids(config.num_instances, config.device), + "joint_ids": make_tensor_joint_ids(config.num_joints, config.device), + } + + +# --- Set Joint Position Target --- +def gen_set_joint_position_target_torch_list(config: BenchmarkConfig) -> dict: + """Generate Torch inputs with list ids for set_joint_position_target.""" + return { + "target": torch.rand(config.num_instances, config.num_joints, device=config.device, dtype=torch.float32), + "env_ids": list(range(config.num_instances)), + "joint_ids": list(range(config.num_joints)), + } + + +def gen_set_joint_position_target_torch_tensor(config: BenchmarkConfig) -> dict: + """Generate Torch inputs with tensor ids for set_joint_position_target.""" + return { + "target": torch.rand(config.num_instances, config.num_joints, device=config.device, dtype=torch.float32), + "env_ids": make_tensor_env_ids(config.num_instances, config.device), + "joint_ids": make_tensor_joint_ids(config.num_joints, config.device), + } + + +# --- Set Joint Velocity Target --- +def gen_set_joint_velocity_target_torch_list(config: BenchmarkConfig) -> dict: + """Generate Torch inputs with list ids for set_joint_velocity_target.""" + return { + "target": torch.rand(config.num_instances, config.num_joints, device=config.device, dtype=torch.float32), + "env_ids": list(range(config.num_instances)), + "joint_ids": list(range(config.num_joints)), + } + + +def gen_set_joint_velocity_target_torch_tensor(config: BenchmarkConfig) -> dict: + """Generate Torch inputs with tensor ids for set_joint_velocity_target.""" + return { + "target": torch.rand(config.num_instances, config.num_joints, device=config.device, dtype=torch.float32), + "env_ids": make_tensor_env_ids(config.num_instances, config.device), + "joint_ids": make_tensor_joint_ids(config.num_joints, config.device), + } + + +# --- Set Joint Effort Target --- +def gen_set_joint_effort_target_torch_list(config: BenchmarkConfig) -> dict: + """Generate Torch inputs with list ids for set_joint_effort_target.""" + return { + "target": torch.rand(config.num_instances, config.num_joints, device=config.device, dtype=torch.float32), + "env_ids": list(range(config.num_instances)), + "joint_ids": list(range(config.num_joints)), + } + + +def gen_set_joint_effort_target_torch_tensor(config: BenchmarkConfig) -> dict: + """Generate Torch inputs with tensor ids for set_joint_effort_target.""" + return { + "target": torch.rand(config.num_instances, config.num_joints, device=config.device, dtype=torch.float32), + "env_ids": make_tensor_env_ids(config.num_instances, config.device), + "joint_ids": make_tensor_joint_ids(config.num_joints, config.device), + } + + +# --- Set Masses --- +def gen_set_masses_torch_list(config: BenchmarkConfig) -> dict: + """Generate Torch inputs with list ids for set_masses.""" + # Articulation masses shape is (N, B) + return { + "masses": torch.rand(config.num_instances, config.num_bodies, device=config.device, dtype=torch.float32), + "env_ids": list(range(config.num_instances)), + "body_ids": list(range(config.num_bodies)), + } + + +def gen_set_masses_torch_tensor(config: BenchmarkConfig) -> dict: + """Generate Torch inputs with tensor ids for set_masses.""" + # Articulation masses shape is (N, B) + return { + "masses": torch.rand(config.num_instances, config.num_bodies, device=config.device, dtype=torch.float32), + "env_ids": make_tensor_env_ids(config.num_instances, config.device), + "body_ids": make_tensor_body_ids(config.num_bodies, config.device), + } + + +# --- Set CoMs --- +def gen_set_coms_torch_list(config: BenchmarkConfig) -> dict: + """Generate Torch inputs with list ids for set_coms.""" + return { + "coms": torch.rand(config.num_instances, config.num_bodies, 7, device=config.device, dtype=torch.float32), + "env_ids": list(range(config.num_instances)), + "body_ids": list(range(config.num_bodies)), + } + + +def gen_set_coms_torch_tensor(config: BenchmarkConfig) -> dict: + """Generate Torch inputs with tensor ids for set_coms.""" + return { + "coms": torch.rand(config.num_instances, config.num_bodies, 7, device=config.device, dtype=torch.float32), + "env_ids": make_tensor_env_ids(config.num_instances, config.device), + "body_ids": make_tensor_body_ids(config.num_bodies, config.device), + } + + +# --- Set Inertias --- +def gen_set_inertias_torch_list(config: BenchmarkConfig) -> dict: + """Generate Torch inputs with list ids for set_inertias.""" + # Articulation inertias shape is (N, B, 9) - flattened 3x3 matrix + return { + "inertias": torch.rand(config.num_instances, config.num_bodies, 9, device=config.device, dtype=torch.float32), + "env_ids": list(range(config.num_instances)), + "body_ids": list(range(config.num_bodies)), + } + + +def gen_set_inertias_torch_tensor(config: BenchmarkConfig) -> dict: + """Generate Torch inputs with tensor ids for set_inertias.""" + # Articulation inertias shape is (N, B, 9) - flattened 3x3 matrix + return { + "inertias": torch.rand(config.num_instances, config.num_bodies, 9, device=config.device, dtype=torch.float32), + "env_ids": make_tensor_env_ids(config.num_instances, config.device), + "body_ids": make_tensor_body_ids(config.num_bodies, config.device), + } + + +# --- Set External Force and Torque --- +def gen_set_external_force_and_torque_torch_list(config: BenchmarkConfig) -> dict: + """Generate Torch inputs with list ids for set_external_force_and_torque.""" + return { + "forces": torch.rand(config.num_instances, config.num_bodies, 3, device=config.device, dtype=torch.float32), + "torques": torch.rand(config.num_instances, config.num_bodies, 3, device=config.device, dtype=torch.float32), + "env_ids": list(range(config.num_instances)), + } + + +def gen_set_external_force_and_torque_torch_tensor(config: BenchmarkConfig) -> dict: + """Generate Torch inputs with tensor ids for set_external_force_and_torque.""" + return { + "forces": torch.rand(config.num_instances, config.num_bodies, 3, device=config.device, dtype=torch.float32), + "torques": torch.rand(config.num_instances, config.num_bodies, 3, device=config.device, dtype=torch.float32), + "env_ids": make_tensor_env_ids(config.num_instances, config.device), + } + + +# ============================================================================= +# Benchmarks +# ============================================================================= + +BENCHMARK_DEPENDENCIES = {} + +BENCHMARKS = [ + # --- Root State (Deprecated) --- + MethodBenchmark( + name="write_root_state_to_sim", + method_name="write_root_state_to_sim", + input_generators={ + "torch_list": gen_root_state_torch_list, + "torch_tensor": gen_root_state_torch_tensor, + }, + category="root_state", + ), + MethodBenchmark( + name="write_root_com_state_to_sim", + method_name="write_root_com_state_to_sim", + input_generators={ + "torch_list": gen_root_com_state_torch_list, + "torch_tensor": gen_root_com_state_torch_tensor, + }, + category="root_state", + ), + MethodBenchmark( + name="write_root_link_state_to_sim", + method_name="write_root_link_state_to_sim", + input_generators={ + "torch_list": gen_root_link_state_torch_list, + "torch_tensor": gen_root_link_state_torch_tensor, + }, + category="root_state", + ), + # --- Root Pose / Velocity --- + MethodBenchmark( + name="write_root_link_pose_to_sim", + method_name="write_root_link_pose_to_sim", + input_generators={ + "torch_list": gen_root_link_pose_torch_list, + "torch_tensor": gen_root_link_pose_torch_tensor, + }, + category="root_pose", + ), + MethodBenchmark( + name="write_root_com_pose_to_sim", + method_name="write_root_com_pose_to_sim", + input_generators={ + "torch_list": gen_root_com_pose_torch_list, + "torch_tensor": gen_root_com_pose_torch_tensor, + }, + category="root_pose", + ), + MethodBenchmark( + name="write_root_link_velocity_to_sim", + method_name="write_root_link_velocity_to_sim", + input_generators={ + "torch_list": gen_root_link_velocity_torch_list, + "torch_tensor": gen_root_link_velocity_torch_tensor, + }, + category="root_velocity", + ), + MethodBenchmark( + name="write_root_com_velocity_to_sim", + method_name="write_root_com_velocity_to_sim", + input_generators={ + "torch_list": gen_root_com_velocity_torch_list, + "torch_tensor": gen_root_com_velocity_torch_tensor, + }, + category="root_velocity", + ), + # --- Joint State --- + MethodBenchmark( + name="write_joint_state_to_sim", + method_name="write_joint_state_to_sim", + input_generators={ + "torch_list": gen_joint_state_torch_list, + "torch_tensor": gen_joint_state_torch_tensor, + }, + category="joint_state", + ), + MethodBenchmark( + name="write_joint_position_to_sim", + method_name="write_joint_position_to_sim", + input_generators={ + "torch_list": gen_joint_position_torch_list, + "torch_tensor": gen_joint_position_torch_tensor, + }, + category="joint_state", + ), + MethodBenchmark( + name="write_joint_velocity_to_sim", + method_name="write_joint_velocity_to_sim", + input_generators={ + "torch_list": gen_joint_velocity_torch_list, + "torch_tensor": gen_joint_velocity_torch_tensor, + }, + category="joint_state", + ), + # --- Joint Params --- + MethodBenchmark( + name="write_joint_stiffness_to_sim", + method_name="write_joint_stiffness_to_sim", + input_generators={ + "torch_list": gen_joint_stiffness_torch_list, + "torch_tensor": gen_joint_stiffness_torch_tensor, + }, + category="joint_params", + ), + MethodBenchmark( + name="write_joint_damping_to_sim", + method_name="write_joint_damping_to_sim", + input_generators={ + "torch_list": gen_joint_damping_torch_list, + "torch_tensor": gen_joint_damping_torch_tensor, + }, + category="joint_params", + ), + MethodBenchmark( + name="write_joint_position_limit_to_sim", + method_name="write_joint_position_limit_to_sim", + input_generators={ + "torch_list": gen_joint_position_limit_torch_list, + "torch_tensor": gen_joint_position_limit_torch_tensor, + }, + category="joint_params", + ), + MethodBenchmark( + name="write_joint_velocity_limit_to_sim", + method_name="write_joint_velocity_limit_to_sim", + input_generators={ + "torch_list": gen_joint_velocity_limit_torch_list, + "torch_tensor": gen_joint_velocity_limit_torch_tensor, + }, + category="joint_params", + ), + MethodBenchmark( + name="write_joint_effort_limit_to_sim", + method_name="write_joint_effort_limit_to_sim", + input_generators={ + "torch_list": gen_joint_effort_limit_torch_list, + "torch_tensor": gen_joint_effort_limit_torch_tensor, + }, + category="joint_params", + ), + MethodBenchmark( + name="write_joint_armature_to_sim", + method_name="write_joint_armature_to_sim", + input_generators={ + "torch_list": gen_joint_armature_torch_list, + "torch_tensor": gen_joint_armature_torch_tensor, + }, + category="joint_params", + ), + MethodBenchmark( + name="write_joint_friction_coefficient_to_sim", + method_name="write_joint_friction_coefficient_to_sim", + input_generators={ + "torch_list": gen_joint_friction_coefficient_torch_list, + "torch_tensor": gen_joint_friction_coefficient_torch_tensor, + }, + category="joint_params", + ), + # --- Joint Targets --- + MethodBenchmark( + name="set_joint_position_target", + method_name="set_joint_position_target", + input_generators={ + "torch_list": gen_set_joint_position_target_torch_list, + "torch_tensor": gen_set_joint_position_target_torch_tensor, + }, + category="joint_targets", + ), + MethodBenchmark( + name="set_joint_velocity_target", + method_name="set_joint_velocity_target", + input_generators={ + "torch_list": gen_set_joint_velocity_target_torch_list, + "torch_tensor": gen_set_joint_velocity_target_torch_tensor, + }, + category="joint_targets", + ), + MethodBenchmark( + name="set_joint_effort_target", + method_name="set_joint_effort_target", + input_generators={ + "torch_list": gen_set_joint_effort_target_torch_list, + "torch_tensor": gen_set_joint_effort_target_torch_tensor, + }, + category="joint_targets", + ), + # --- Body Properties --- + MethodBenchmark( + name="set_masses", + method_name="set_masses", + input_generators={ + "torch_list": gen_set_masses_torch_list, + "torch_tensor": gen_set_masses_torch_tensor, + }, + category="body_props", + ), + MethodBenchmark( + name="set_coms", + method_name="set_coms", + input_generators={ + "torch_list": gen_set_coms_torch_list, + "torch_tensor": gen_set_coms_torch_tensor, + }, + category="body_props", + ), + MethodBenchmark( + name="set_inertias", + method_name="set_inertias", + input_generators={ + "torch_list": gen_set_inertias_torch_list, + "torch_tensor": gen_set_inertias_torch_tensor, + }, + category="body_props", + ), + MethodBenchmark( + name="set_external_force_and_torque", + method_name="set_external_force_and_torque", + input_generators={ + "torch_list": gen_set_external_force_and_torque_torch_list, + "torch_tensor": gen_set_external_force_and_torque_torch_tensor, + }, + category="external_wrench", + ), +] + + +def run_benchmark(config: BenchmarkConfig): + """Run all benchmarks.""" + results = [] + + # Check if we should run all modes or specific ones + modes_to_run = [] + if isinstance(config.mode, str): + if config.mode == "all": + # Will be populated dynamically based on available generators + modes_to_run = None + else: + modes_to_run = [config.mode] + elif isinstance(config.mode, list): + modes_to_run = config.mode + + # Setup + articulation, _, _ = create_test_articulation( + num_instances=config.num_instances, + num_bodies=config.num_bodies, + num_joints=config.num_joints, + device=config.device, + ) + + print( + f"Benchmarking Articulation (PhysX) with {config.num_instances} instances, {config.num_bodies} bodies," + f" {config.num_joints} joints..." + ) + print(f"Device: {config.device}") + print(f"Iterations: {config.num_iterations}, Warmup: {config.warmup_steps}") + print(f"Modes: {modes_to_run if modes_to_run else 'All available'}") + + print(f"\nBenchmarking {len(BENCHMARKS)} methods...") + for i, benchmark in enumerate(BENCHMARKS): + method = getattr(articulation, benchmark.method_name, None) + + # Determine which modes to run for this benchmark + available_modes = list(benchmark.input_generators.keys()) + current_modes = modes_to_run if modes_to_run is not None else available_modes + + # Filter modes that are available for this benchmark + current_modes = [m for m in current_modes if m in available_modes] + + for mode in current_modes: + generator = benchmark.input_generators[mode] + print(f"[{i + 1}/{len(BENCHMARKS)}] [{mode.upper()}] {benchmark.name}...", end=" ", flush=True) + + result = benchmark_method( + method=method, + method_name=benchmark.name, + generator=generator, + config=config, + dependencies=BENCHMARK_DEPENDENCIES, + ) + result.mode = mode + results.append(result) + + if result.skipped: + print(f"SKIPPED ({result.skip_reason})") + else: + print(f"{result.mean_time_us:.2f} ± {result.std_time_us:.2f} µs") + + return results + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Benchmark Articulation methods (PhysX backend).") + parser.add_argument("--num_iterations", type=int, default=1000, help="Number of iterations") + parser.add_argument("--warmup_steps", type=int, default=10, help="Number of warmup steps") + parser.add_argument("--num_instances", type=int, default=4096, help="Number of instances") + parser.add_argument("--num_bodies", type=int, default=12, help="Number of bodies") + parser.add_argument("--num_joints", type=int, default=11, help="Number of joints") + parser.add_argument("--device", type=str, default="cuda:0", help="Device") + parser.add_argument("--mode", type=str, default="all", help="Benchmark mode (all, torch_list, torch_tensor)") + parser.add_argument("--output", type=str, default=None, help="Output JSON filename") + parser.add_argument("--no_csv", action="store_true", help="Disable CSV output") + + args = parser.parse_args() + + config = BenchmarkConfig( + num_iterations=args.num_iterations, + warmup_steps=args.warmup_steps, + num_instances=args.num_instances, + num_bodies=args.num_bodies, + num_joints=args.num_joints, + device=args.device, + mode=args.mode, + ) + + results = run_benchmark(config) + + hardware_info = get_hardware_info() + print_hardware_info(hardware_info) + print_results(results) + + if args.output: + json_filename = args.output + else: + json_filename = get_default_output_filename("articulation_benchmark") + + export_results_json(results, config, hardware_info, json_filename) + + if not args.no_csv: + csv_filename = json_filename.replace(".json", ".csv") + export_results_csv(results, csv_filename) diff --git a/source/isaaclab_physx/benchmark/assets/benchmark_articulation_data.py b/source/isaaclab_physx/benchmark/assets/benchmark_articulation_data.py new file mode 100644 index 00000000000..ae1ae653c51 --- /dev/null +++ b/source/isaaclab_physx/benchmark/assets/benchmark_articulation_data.py @@ -0,0 +1,420 @@ +# 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 + +"""Micro-benchmarking framework for ArticulationData class. + +This module provides a benchmarking framework to measure the performance of all functions +in the ArticulationData class. Each function is run multiple times with randomized mock data, +and timing statistics (mean and standard deviation) are reported. + +Usage: + python benchmark_articulation_data.py [--num_iterations N] [--warmup_steps W] + [--num_instances I] [--num_bodies B] [--num_joints J] + +Example: + python benchmark_articulation_data.py --num_iterations 10000 --warmup_steps 10 +""" + +from __future__ import annotations + +import argparse +import sys +import warnings +from types import ModuleType + +import torch +import warp as wp + +# Initialize Warp first +wp.init() + + +# ============================================================================= +# Mock Setup - Must happen BEFORE importing ArticulationData +# ============================================================================= + + +class MockPhysicsSimView: + """Simple mock for the physics simulation view.""" + + def get_gravity(self): + """Return gravity as a tuple of 3 floats.""" + return (0.0, 0.0, -9.81) + + def update_articulations_kinematic(self): + """No-op for kinematic updates.""" + pass + + +class MockSimulationManager: + """Simple mock for SimulationManager.""" + + @staticmethod + def get_physics_sim_view(): + return MockPhysicsSimView() + + +# Mock isaacsim.core.simulation_manager +mock_sim_manager_module = ModuleType("isaacsim.core.simulation_manager") +mock_sim_manager_module.SimulationManager = MockSimulationManager +sys.modules["isaacsim"] = ModuleType("isaacsim") +sys.modules["isaacsim.core"] = ModuleType("isaacsim.core") +sys.modules["isaacsim.core.simulation_manager"] = mock_sim_manager_module + + +# Mock BaseArticulationData - this is just an abstract class, so we provide an empty one +# to avoid importing the whole isaaclab.assets package +class BaseArticulationData: + """Mock base class to avoid importing isaaclab.assets (which has many dependencies).""" + + def __init__(self, root_view, device: str): + self.device = device + + +# Create mock module for isaaclab.assets.articulation.base_articulation_data +mock_base_module = ModuleType("isaaclab.assets.articulation.base_articulation_data") +mock_base_module.BaseArticulationData = BaseArticulationData +sys.modules["isaaclab.assets.articulation.base_articulation_data"] = mock_base_module + +# Mock pxr (USD library - not available in headless docker, used by isaaclab.utils.mesh) +from unittest.mock import MagicMock + +sys.modules["pxr"] = MagicMock() +sys.modules["pxr.Usd"] = MagicMock() +sys.modules["pxr.UsdGeom"] = MagicMock() + +# Now we can directly import ArticulationData +import importlib.util +from pathlib import Path + +benchmark_dir = Path(__file__).resolve().parent +articulation_data_path = ( + benchmark_dir.parents[1] / "isaaclab_physx" / "assets" / "articulation" / "articulation_data.py" +) + +spec = importlib.util.spec_from_file_location( + "isaaclab_physx.assets.articulation.articulation_data", articulation_data_path +) +articulation_data_module = importlib.util.module_from_spec(spec) +sys.modules["isaaclab_physx.assets.articulation.articulation_data"] = articulation_data_module +spec.loader.exec_module(articulation_data_module) +ArticulationData = articulation_data_module.ArticulationData + +# Import shared utilities from common module +# Import mock classes from PhysX test utilities +from isaaclab_physx.test.mock_interfaces.views import MockArticulationView + +from isaaclab.test.benchmark import ( + BenchmarkConfig, + BenchmarkResult, + MethodBenchmark, + benchmark_method, + export_results_csv, + export_results_json, + get_default_output_filename, + get_hardware_info, + print_hardware_info, + print_results, +) + +# Suppress deprecation warnings during benchmarking +warnings.filterwarnings("ignore", category=DeprecationWarning) +warnings.filterwarnings("ignore", category=UserWarning) + + +# ============================================================================= +# Skip Lists +# ============================================================================= + +# List of deprecated properties (for backward compatibility) - skip these +DEPRECATED_PROPERTIES = { + "default_root_state", + "root_pose_w", + "root_pos_w", + "root_quat_w", + "root_vel_w", + "root_lin_vel_w", + "root_ang_vel_w", + "root_lin_vel_b", + "root_ang_vel_b", + "body_pose_w", + "body_pos_w", + "body_quat_w", + "body_vel_w", + "body_lin_vel_w", + "body_ang_vel_w", + "body_acc_w", + "body_lin_acc_w", + "body_ang_acc_w", + "com_pos_b", + "com_quat_b", + "joint_limits", + "joint_friction", + "fixed_tendon_limit", + "applied_torque", + "computed_torque", + "joint_dynamic_friction", + "joint_effort_target", + "joint_viscous_friction", + "joint_velocity_limits", + # Also skip the combined state properties marked as deprecated + "root_state_w", + "root_link_state_w", + "root_com_state_w", + "body_state_w", + "body_link_state_w", + "body_com_state_w", +} + +# List of properties that raise NotImplementedError - skip these +NOT_IMPLEMENTED_PROPERTIES = { + "fixed_tendon_stiffness", + "fixed_tendon_damping", + "fixed_tendon_limit_stiffness", + "fixed_tendon_rest_length", + "fixed_tendon_offset", + "fixed_tendon_pos_limits", + "spatial_tendon_stiffness", + "spatial_tendon_damping", + "spatial_tendon_limit_stiffness", + "spatial_tendon_offset", + "body_incoming_joint_wrench_b", +} + +# Removed default_* properties that raise RuntimeError +REMOVED_PROPERTIES = { + "default_fixed_tendon_damping", + "default_fixed_tendon_limit", + "default_fixed_tendon_limit_stiffness", + "default_fixed_tendon_offset", + "default_fixed_tendon_pos_limits", + "default_fixed_tendon_rest_length", + "default_fixed_tendon_stiffness", + "default_inertia", + "default_joint_armature", + "default_joint_damping", + "default_joint_dynamic_friction_coeff", + "default_joint_friction", + "default_joint_friction_coeff", + "default_joint_limits", + "default_joint_pos_limits", + "default_joint_stiffness", + "default_joint_viscous_friction_coeff", + "default_mass", + "default_spatial_tendon_damping", + "default_spatial_tendon_limit_stiffness", + "default_spatial_tendon_offset", + "default_spatial_tendon_stiffness", +} + +# Private/internal properties and methods to skip +INTERNAL_PROPERTIES = { + "_create_simulation_bindings", + "_create_buffers", + "update", + "is_primed", + "device", + "body_names", + "joint_names", + "fixed_tendon_names", + "spatial_tendon_names", + "GRAVITY_VEC_W", + "GRAVITY_VEC_W_TORCH", + "FORWARD_VEC_B", + "FORWARD_VEC_B_TORCH", + "ALL_ENV_MASK", + "ALL_BODY_MASK", + "ALL_JOINT_MASK", + "ENV_MASK", + "BODY_MASK", + "JOINT_MASK", +} + +# Dependency mapping for properties +PROPERTY_DEPENDENCIES = { + "root_link_lin_vel_w": ["root_link_vel_w"], + "root_link_ang_vel_w": ["root_link_vel_w"], + "root_link_lin_vel_b": ["root_link_vel_b"], + "root_link_ang_vel_b": ["root_link_vel_b"], + "root_com_pos_w": ["root_com_pose_w"], + "root_com_quat_w": ["root_com_pose_w"], + "root_com_lin_vel_b": ["root_com_vel_b"], + "root_com_ang_vel_b": ["root_com_vel_b"], + "root_com_lin_vel_w": ["root_com_vel_w"], + "root_com_ang_vel_w": ["root_com_vel_w"], + "root_link_pos_w": ["root_link_pose_w"], + "root_link_quat_w": ["root_link_pose_w"], + "body_link_lin_vel_w": ["body_link_vel_w"], + "body_link_ang_vel_w": ["body_link_vel_w"], + "body_link_pos_w": ["body_link_pose_w"], + "body_link_quat_w": ["body_link_pose_w"], + "body_com_pos_w": ["body_com_pose_w"], + "body_com_quat_w": ["body_com_pose_w"], + "body_com_lin_vel_w": ["body_com_vel_w"], + "body_com_ang_vel_w": ["body_com_vel_w"], + "body_com_lin_acc_w": ["body_com_acc_w"], + "body_com_ang_acc_w": ["body_com_acc_w"], + "body_com_quat_b": ["body_com_pose_b"], +} + + +# ============================================================================= +# Benchmark Functions +# ============================================================================= + + +def get_benchmarkable_properties(articulation_data: ArticulationData) -> list[str]: + """Get list of properties that can be benchmarked.""" + all_properties = [] + + for name in dir(articulation_data): + if name.startswith("_"): + continue + if name in DEPRECATED_PROPERTIES: + continue + if name in NOT_IMPLEMENTED_PROPERTIES: + continue + if name in REMOVED_PROPERTIES: + continue + if name in INTERNAL_PROPERTIES: + continue + + try: + attr = getattr(type(articulation_data), name, None) + if isinstance(attr, property): + all_properties.append(name) + except Exception: + pass + + return sorted(all_properties) + + +def setup_mock_environment(config: BenchmarkConfig) -> MockArticulationView: + """Set up the mock environment for benchmarking.""" + mock_view = MockArticulationView( + count=config.num_instances, + num_links=config.num_bodies, + num_dofs=config.num_joints, + device=config.device, + ) + return mock_view + + +def run_benchmarks(config: BenchmarkConfig) -> list[BenchmarkResult]: + """Run all benchmarks for ArticulationData.""" + results = [] + + # Setup mock environment + mock_view = setup_mock_environment(config) + mock_view.set_random_mock_data() + + # Create ArticulationData instance + articulation_data = ArticulationData(mock_view, config.device) + + # Get list of properties to benchmark + properties = get_benchmarkable_properties(articulation_data) + + # Generator that updates mock data and invalidates timestamp + def gen_mock_data(cfg: BenchmarkConfig) -> dict: + mock_view.set_mock_coms(torch.randn(cfg.num_instances, cfg.num_bodies, 7, device=cfg.device)) + mock_view.set_mock_inertias(torch.randn(cfg.num_instances, cfg.num_bodies, 3, 3, device=cfg.device)) + mock_view.set_mock_root_transforms(torch.randn(cfg.num_instances, 7, device=cfg.device)) + mock_view.set_mock_root_velocities(torch.randn(cfg.num_instances, 6, device=cfg.device)) + mock_view.set_mock_link_transforms(torch.randn(cfg.num_instances, cfg.num_bodies, 7, device=cfg.device)) + mock_view.set_mock_link_velocities(torch.randn(cfg.num_instances, cfg.num_bodies, 6, device=cfg.device)) + mock_view.set_mock_link_accelerations(torch.randn(cfg.num_instances, cfg.num_bodies, 6, device=cfg.device)) + mock_view.set_mock_dof_positions(torch.randn(cfg.num_instances, cfg.num_joints, device=cfg.device)) + mock_view.set_mock_dof_velocities(torch.randn(cfg.num_instances, cfg.num_joints, device=cfg.device)) + articulation_data._sim_timestamp += 1.0 + return {} + + # Create benchmarks dynamically + benchmarks = [] + for prop_name in properties: + benchmarks.append( + MethodBenchmark( + name=prop_name, + method_name=prop_name, + input_generators={"default": gen_mock_data}, + category="property", + ) + ) + + print(f"\nBenchmarking {len(benchmarks)} properties...") + print(f"Config: {config.num_iterations} iterations, {config.warmup_steps} warmup steps") + print(f" {config.num_instances} instances, {config.num_bodies} bodies, {config.num_joints} joints") + print("-" * 80) + + for i, benchmark in enumerate(benchmarks): + + def prop_accessor(prop=benchmark.method_name, **kwargs): + return getattr(articulation_data, prop) + + print(f"[{i + 1}/{len(benchmarks)}] [DEFAULT] {benchmark.name}...", end=" ", flush=True) + + result = benchmark_method( + method=prop_accessor, + method_name=benchmark.name, + generator=gen_mock_data, + config=config, + dependencies=PROPERTY_DEPENDENCIES, + ) + result.mode = "default" + results.append(result) + + if result.skipped: + print(f"SKIPPED ({result.skip_reason})") + else: + print(f"{result.mean_time_us:.2f} ± {result.std_time_us:.2f} µs") + + return results + + +def main(): + """Main entry point for the benchmarking script.""" + parser = argparse.ArgumentParser( + description="Micro-benchmarking framework for ArticulationData class.", + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + ) + parser.add_argument("--num_iterations", type=int, default=1000, help="Number of iterations") + parser.add_argument("--warmup_steps", type=int, default=10, help="Number of warmup steps") + parser.add_argument("--num_instances", type=int, default=4096, help="Number of instances") + parser.add_argument("--num_bodies", type=int, default=12, help="Number of bodies") + parser.add_argument("--num_joints", type=int, default=11, help="Number of joints") + parser.add_argument("--device", type=str, default="cuda:0", help="Device") + parser.add_argument("--output", type=str, default=None, help="Output JSON filename") + parser.add_argument("--no_csv", action="store_true", help="Disable CSV output") + + args = parser.parse_args() + + config = BenchmarkConfig( + num_iterations=args.num_iterations, + warmup_steps=args.warmup_steps, + num_instances=args.num_instances, + num_bodies=args.num_bodies, + num_joints=args.num_joints, + device=args.device, + ) + + results = run_benchmarks(config) + + hardware_info = get_hardware_info() + print_hardware_info(hardware_info) + print_results(results, include_mode=False) + + if args.output: + json_filename = args.output + else: + json_filename = get_default_output_filename("articulation_data_benchmark") + + export_results_json(results, config, hardware_info, json_filename) + + if not args.no_csv: + csv_filename = json_filename.replace(".json", ".csv") + export_results_csv(results, csv_filename) + + +if __name__ == "__main__": + main() diff --git a/source/isaaclab_physx/benchmark/assets/benchmark_rigid_object.py b/source/isaaclab_physx/benchmark/assets/benchmark_rigid_object.py new file mode 100644 index 00000000000..ce48ab12d51 --- /dev/null +++ b/source/isaaclab_physx/benchmark/assets/benchmark_rigid_object.py @@ -0,0 +1,704 @@ +# 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 + +"""Micro-benchmarking framework for RigidObject class (PhysX backend). + +This module provides a benchmarking framework to measure the performance of setter and writer +methods in the RigidObject class. Each method is benchmarked under two scenarios: + +1. **Torch List**: Inputs are PyTorch tensors with list indices. +2. **Torch Tensor**: Inputs are PyTorch tensors with tensor indices. + +Usage: + python benchmark_rigid_object.py [--num_iterations N] [--warmup_steps W] [--num_instances I] + +Example: + python benchmark_rigid_object.py --num_iterations 1000 --warmup_steps 10 + python benchmark_rigid_object.py --mode torch_list # Only run list-based benchmarks + python benchmark_rigid_object.py --mode torch_tensor # Only run tensor-based benchmarks +""" + +from __future__ import annotations + +import argparse +import sys +import warnings +from types import ModuleType +from unittest.mock import MagicMock + +import torch +import warp as wp + +# Initialize Warp +wp.init() + + +# ============================================================================= +# Mock Setup - Must happen BEFORE importing RigidObject +# ============================================================================= + + +class MockPhysicsSimView: + """Simple mock for the physics simulation view.""" + + def get_gravity(self): + return (0.0, 0.0, -9.81) + + +class MockSimulationManager: + """Simple mock for SimulationManager.""" + + @staticmethod + def get_physics_sim_view(): + return MockPhysicsSimView() + + +# Mock isaacsim.core.simulation_manager +mock_sim_manager_module = ModuleType("isaacsim.core.simulation_manager") +mock_sim_manager_module.SimulationManager = MockSimulationManager +sys.modules["isaacsim"] = ModuleType("isaacsim") +sys.modules["isaacsim.core"] = ModuleType("isaacsim.core") +sys.modules["isaacsim.core.simulation_manager"] = mock_sim_manager_module + +# Mock pxr (USD library) +sys.modules["pxr"] = MagicMock() +sys.modules["pxr.Usd"] = MagicMock() +sys.modules["pxr.UsdGeom"] = MagicMock() +sys.modules["pxr.UsdPhysics"] = MagicMock() + +# Mock omni module hierarchy (must be ModuleType for proper package behavior) +omni_mocks = [ + "omni", + "omni.kit", + "omni.kit.app", + "omni.kit.commands", + "omni.usd", + "omni.client", + "omni.timeline", + "omni.physx", + "omni.physx.scripts", + "omni.physx.scripts.utils", + "omni.physics", + "omni.physics.tensors", + "omni.physics.tensors.impl", + "omni.physics.tensors.impl.api", +] +for mod_name in omni_mocks: + mock = MagicMock() + mock.__name__ = mod_name + mock.__path__ = [] + mock.__package__ = mod_name + sys.modules[mod_name] = mock + +# Mock carb (needed by isaaclab.utils.assets) +mock_carb = MagicMock() +mock_carb.settings.get_settings.return_value.get.return_value = "/mock/path" +sys.modules["carb"] = mock_carb + +# Mock isaaclab.sim module hierarchy (to avoid importing converters) +sim_mock = MagicMock() +sim_mock.find_first_matching_prim = MagicMock() +sim_mock.get_all_matching_child_prims = MagicMock(return_value=[]) +sys.modules["isaaclab.sim"] = sim_mock +sys.modules["isaaclab.sim.utils"] = MagicMock() +sys.modules["isaaclab.sim.utils.stage"] = MagicMock() +sys.modules["isaaclab.sim.converters"] = MagicMock() + +# Mock WrenchComposer - import from mock_interfaces and patch into the module +from isaaclab.test.mock_interfaces.utils import MockWrenchComposer + +mock_wrench_module = ModuleType("isaaclab.utils.wrench_composer") +mock_wrench_module.WrenchComposer = MockWrenchComposer +sys.modules["isaaclab.utils.wrench_composer"] = mock_wrench_module + + +# Mock base classes to avoid importing full isaaclab.assets package +class BaseRigidObject: + """Mock base class.""" + + @property + def device(self) -> str: + return self._device + + +class BaseRigidObjectData: + """Mock base class.""" + + def __init__(self, root_view, device: str): + self.device = device + + +mock_base_rigid_object = ModuleType("isaaclab.assets.rigid_object.base_rigid_object") +mock_base_rigid_object.BaseRigidObject = BaseRigidObject +sys.modules["isaaclab.assets.rigid_object.base_rigid_object"] = mock_base_rigid_object + +mock_base_rigid_object_data = ModuleType("isaaclab.assets.rigid_object.base_rigid_object_data") +mock_base_rigid_object_data.BaseRigidObjectData = BaseRigidObjectData +sys.modules["isaaclab.assets.rigid_object.base_rigid_object_data"] = mock_base_rigid_object_data + +# Mock RigidObjectCfg +mock_cfg_module = ModuleType("isaaclab.assets.rigid_object.rigid_object_cfg") +mock_cfg_module.RigidObjectCfg = type("RigidObjectCfg", (), {"prim_path": None}) +sys.modules["isaaclab.assets.rigid_object.rigid_object_cfg"] = mock_cfg_module + +# Now import via importlib to bypass __init__.py +import importlib.util +from pathlib import Path + +benchmark_dir = Path(__file__).resolve().parent + +# Load RigidObjectData +rigid_object_data_path = ( + benchmark_dir.parents[1] / "isaaclab_physx" / "assets" / "rigid_object" / "rigid_object_data.py" +) +spec = importlib.util.spec_from_file_location( + "isaaclab_physx.assets.rigid_object.rigid_object_data", rigid_object_data_path +) +rigid_object_data_module = importlib.util.module_from_spec(spec) +sys.modules["isaaclab_physx.assets.rigid_object.rigid_object_data"] = rigid_object_data_module +spec.loader.exec_module(rigid_object_data_module) +RigidObjectData = rigid_object_data_module.RigidObjectData + +# Load RigidObject +rigid_object_path = benchmark_dir.parents[1] / "isaaclab_physx" / "assets" / "rigid_object" / "rigid_object.py" +spec = importlib.util.spec_from_file_location("isaaclab_physx.assets.rigid_object.rigid_object", rigid_object_path) +rigid_object_module = importlib.util.module_from_spec(spec) +sys.modules["isaaclab_physx.assets.rigid_object.rigid_object"] = rigid_object_module +spec.loader.exec_module(rigid_object_module) +RigidObject = rigid_object_module.RigidObject + + +# Simple RigidObjectCfg for testing +class RigidObjectCfg: + def __init__(self, prim_path: str = "/World/Object"): + self.prim_path = prim_path + + +# Import shared utilities from common module +# Import mock classes from PhysX test utilities +from isaaclab_physx.test.mock_interfaces.views import MockRigidBodyView + +from isaaclab.test.benchmark import ( + BenchmarkConfig, + MethodBenchmark, + benchmark_method, + export_results_csv, + export_results_json, + get_default_output_filename, + get_hardware_info, + make_tensor_env_ids, + print_hardware_info, + print_results, +) + +# Suppress deprecation warnings during benchmarking +warnings.filterwarnings("ignore", category=DeprecationWarning) +warnings.filterwarnings("ignore", category=UserWarning) + +# Suppress isaaclab logging (deprecation warnings) +import logging + +logging.getLogger("isaaclab_physx").setLevel(logging.ERROR) +logging.getLogger("isaaclab").setLevel(logging.ERROR) + + +def create_test_rigid_object( + num_instances: int = 2, + num_bodies: int = 1, + device: str = "cuda:0", +) -> tuple[RigidObject, MockRigidBodyView, MagicMock]: + """Create a test RigidObject instance with mocked dependencies.""" + rigid_object = object.__new__(RigidObject) + + rigid_object.cfg = RigidObjectCfg( + prim_path="/World/Object", + ) + + # Create PhysX mock view + mock_view = MockRigidBodyView( + count=num_instances, + device=device, + ) + mock_view.set_random_mock_data() + + # Set up attributes required before _create_buffers + object.__setattr__(rigid_object, "_root_view", mock_view) + object.__setattr__(rigid_object, "_device", device) + + # Create RigidObjectData instance (mocks already set up at module level) + data = RigidObjectData(mock_view, device) + object.__setattr__(rigid_object, "_data", data) + + # Call _create_buffers to set up all internal buffers and wrench composers + rigid_object._create_buffers() + + return rigid_object, mock_view, None + + +# ============================================================================= +# Input Generators (Torch-only for PhysX backend) +# ============================================================================= + + +# --- Root State (Deprecated) --- +def gen_root_state_torch_list(config: BenchmarkConfig) -> dict: + """Generate Torch inputs with list env_ids for write_root_state_to_sim.""" + return { + "root_state": torch.rand(config.num_instances, 13, device=config.device, dtype=torch.float32), + "env_ids": list(range(config.num_instances)), + } + + +def gen_root_state_torch_tensor(config: BenchmarkConfig) -> dict: + """Generate Torch inputs with tensor env_ids for write_root_state_to_sim.""" + return { + "root_state": torch.rand(config.num_instances, 13, device=config.device, dtype=torch.float32), + "env_ids": make_tensor_env_ids(config.num_instances, config.device), + } + + +# --- Root COM State (Deprecated) --- +def gen_root_com_state_torch_list(config: BenchmarkConfig) -> dict: + """Generate Torch inputs with list env_ids for write_root_com_state_to_sim.""" + return { + "root_state": torch.rand(config.num_instances, 13, device=config.device, dtype=torch.float32), + "env_ids": list(range(config.num_instances)), + } + + +def gen_root_com_state_torch_tensor(config: BenchmarkConfig) -> dict: + """Generate Torch inputs with tensor env_ids for write_root_com_state_to_sim.""" + return { + "root_state": torch.rand(config.num_instances, 13, device=config.device, dtype=torch.float32), + "env_ids": make_tensor_env_ids(config.num_instances, config.device), + } + + +# --- Root Link State (Deprecated) --- +def gen_root_link_state_torch_list(config: BenchmarkConfig) -> dict: + """Generate Torch inputs with list env_ids for write_root_link_state_to_sim.""" + return { + "root_state": torch.rand(config.num_instances, 13, device=config.device, dtype=torch.float32), + "env_ids": list(range(config.num_instances)), + } + + +def gen_root_link_state_torch_tensor(config: BenchmarkConfig) -> dict: + """Generate Torch inputs with tensor env_ids for write_root_link_state_to_sim.""" + return { + "root_state": torch.rand(config.num_instances, 13, device=config.device, dtype=torch.float32), + "env_ids": make_tensor_env_ids(config.num_instances, config.device), + } + + +# --- Root Pose (Deprecated) --- +def gen_root_pose_torch_list(config: BenchmarkConfig) -> dict: + """Generate Torch inputs with list env_ids for write_root_pose_to_sim.""" + return { + "root_pose": torch.rand(config.num_instances, 7, device=config.device, dtype=torch.float32), + "env_ids": list(range(config.num_instances)), + } + + +def gen_root_pose_torch_tensor(config: BenchmarkConfig) -> dict: + """Generate Torch inputs with tensor env_ids for write_root_pose_to_sim.""" + return { + "root_pose": torch.rand(config.num_instances, 7, device=config.device, dtype=torch.float32), + "env_ids": make_tensor_env_ids(config.num_instances, config.device), + } + + +# --- Root Link Pose --- +def gen_root_link_pose_torch_list(config: BenchmarkConfig) -> dict: + """Generate Torch inputs with list env_ids for write_root_link_pose_to_sim.""" + return { + "root_pose": torch.rand(config.num_instances, 7, device=config.device, dtype=torch.float32), + "env_ids": list(range(config.num_instances)), + } + + +def gen_root_link_pose_torch_tensor(config: BenchmarkConfig) -> dict: + """Generate Torch inputs with tensor env_ids for write_root_link_pose_to_sim.""" + return { + "root_pose": torch.rand(config.num_instances, 7, device=config.device, dtype=torch.float32), + "env_ids": make_tensor_env_ids(config.num_instances, config.device), + } + + +# --- Root COM Pose --- +def gen_root_com_pose_torch_list(config: BenchmarkConfig) -> dict: + """Generate Torch inputs with list env_ids for write_root_com_pose_to_sim.""" + return { + "root_pose": torch.rand(config.num_instances, 7, device=config.device, dtype=torch.float32), + "env_ids": list(range(config.num_instances)), + } + + +def gen_root_com_pose_torch_tensor(config: BenchmarkConfig) -> dict: + """Generate Torch inputs with tensor env_ids for write_root_com_pose_to_sim.""" + return { + "root_pose": torch.rand(config.num_instances, 7, device=config.device, dtype=torch.float32), + "env_ids": make_tensor_env_ids(config.num_instances, config.device), + } + + +# --- Root Velocity (Deprecated) --- +def gen_root_velocity_torch_list(config: BenchmarkConfig) -> dict: + """Generate Torch inputs with list env_ids for write_root_velocity_to_sim.""" + return { + "root_velocity": torch.rand(config.num_instances, 6, device=config.device, dtype=torch.float32), + "env_ids": list(range(config.num_instances)), + } + + +def gen_root_velocity_torch_tensor(config: BenchmarkConfig) -> dict: + """Generate Torch inputs with tensor env_ids for write_root_velocity_to_sim.""" + return { + "root_velocity": torch.rand(config.num_instances, 6, device=config.device, dtype=torch.float32), + "env_ids": make_tensor_env_ids(config.num_instances, config.device), + } + + +# --- Root Link Velocity --- +def gen_root_link_velocity_torch_list(config: BenchmarkConfig) -> dict: + """Generate Torch inputs with list env_ids for write_root_link_velocity_to_sim.""" + return { + "root_velocity": torch.rand(config.num_instances, 6, device=config.device, dtype=torch.float32), + "env_ids": list(range(config.num_instances)), + } + + +def gen_root_link_velocity_torch_tensor(config: BenchmarkConfig) -> dict: + """Generate Torch inputs with tensor env_ids for write_root_link_velocity_to_sim.""" + return { + "root_velocity": torch.rand(config.num_instances, 6, device=config.device, dtype=torch.float32), + "env_ids": make_tensor_env_ids(config.num_instances, config.device), + } + + +# --- Root COM Velocity --- +def gen_root_com_velocity_torch_list(config: BenchmarkConfig) -> dict: + """Generate Torch inputs with list env_ids for write_root_com_velocity_to_sim.""" + return { + "root_velocity": torch.rand(config.num_instances, 6, device=config.device, dtype=torch.float32), + "env_ids": list(range(config.num_instances)), + } + + +def gen_root_com_velocity_torch_tensor(config: BenchmarkConfig) -> dict: + """Generate Torch inputs with tensor env_ids for write_root_com_velocity_to_sim.""" + return { + "root_velocity": torch.rand(config.num_instances, 6, device=config.device, dtype=torch.float32), + "env_ids": make_tensor_env_ids(config.num_instances, config.device), + } + + +# --- Masses --- +def gen_masses_torch_list(config: BenchmarkConfig) -> dict: + """Generate Torch inputs with list ids for set_masses.""" + # RigidObject has only 1 body, don't pass body_ids to avoid advanced indexing issues + return { + "masses": torch.rand(config.num_instances, config.num_bodies, device=config.device, dtype=torch.float32), + "env_ids": list(range(config.num_instances)), + } + + +def gen_masses_torch_tensor(config: BenchmarkConfig) -> dict: + """Generate Torch inputs with tensor ids for set_masses.""" + return { + "masses": torch.rand(config.num_instances, config.num_bodies, device=config.device, dtype=torch.float32), + "env_ids": make_tensor_env_ids(config.num_instances, config.device), + } + + +# --- CoMs --- +def gen_coms_torch_list(config: BenchmarkConfig) -> dict: + """Generate Torch inputs with list ids for set_coms.""" + # RigidObject has only 1 body, don't pass body_ids to avoid advanced indexing issues + return { + "coms": torch.rand(config.num_instances, config.num_bodies, 3, device=config.device, dtype=torch.float32), + "env_ids": list(range(config.num_instances)), + } + + +def gen_coms_torch_tensor(config: BenchmarkConfig) -> dict: + """Generate Torch inputs with tensor ids for set_coms.""" + return { + "coms": torch.rand(config.num_instances, config.num_bodies, 3, device=config.device, dtype=torch.float32), + "env_ids": make_tensor_env_ids(config.num_instances, config.device), + } + + +# --- Inertias --- +def gen_inertias_torch_list(config: BenchmarkConfig) -> dict: + """Generate Torch inputs with list ids for set_inertias.""" + # RigidObject has only 1 body, don't pass body_ids to avoid advanced indexing issues + return { + "inertias": torch.rand( + config.num_instances, config.num_bodies, 3, 3, device=config.device, dtype=torch.float32 + ), + "env_ids": list(range(config.num_instances)), + } + + +def gen_inertias_torch_tensor(config: BenchmarkConfig) -> dict: + """Generate Torch inputs with tensor ids for set_inertias.""" + return { + "inertias": torch.rand( + config.num_instances, config.num_bodies, 3, 3, device=config.device, dtype=torch.float32 + ), + "env_ids": make_tensor_env_ids(config.num_instances, config.device), + } + + +# --- External Wrench --- +def gen_external_force_and_torque_torch_list(config: BenchmarkConfig) -> dict: + """Generate Torch inputs with list ids for set_external_force_and_torque.""" + # Note: body_ids is not used by set_external_force_and_torque (it uses internal _ALL_BODY_INDICES_WP) + return { + "forces": torch.rand(config.num_instances, config.num_bodies, 3, device=config.device, dtype=torch.float32), + "torques": torch.rand(config.num_instances, config.num_bodies, 3, device=config.device, dtype=torch.float32), + "env_ids": list(range(config.num_instances)), + } + + +def gen_external_force_and_torque_torch_tensor(config: BenchmarkConfig) -> dict: + """Generate Torch inputs with tensor ids for set_external_force_and_torque.""" + return { + "forces": torch.rand(config.num_instances, config.num_bodies, 3, device=config.device, dtype=torch.float32), + "torques": torch.rand(config.num_instances, config.num_bodies, 3, device=config.device, dtype=torch.float32), + "env_ids": make_tensor_env_ids(config.num_instances, config.device), + } + + +# ============================================================================= +# Benchmarks +# ============================================================================= + +BENCHMARK_DEPENDENCIES = {} + +BENCHMARKS = [ + # --- Root State (Deprecated) --- + MethodBenchmark( + name="write_root_state_to_sim", + method_name="write_root_state_to_sim", + input_generators={ + "torch_list": gen_root_state_torch_list, + "torch_tensor": gen_root_state_torch_tensor, + }, + category="root_state", + ), + MethodBenchmark( + name="write_root_com_state_to_sim", + method_name="write_root_com_state_to_sim", + input_generators={ + "torch_list": gen_root_com_state_torch_list, + "torch_tensor": gen_root_com_state_torch_tensor, + }, + category="root_state", + ), + MethodBenchmark( + name="write_root_link_state_to_sim", + method_name="write_root_link_state_to_sim", + input_generators={ + "torch_list": gen_root_link_state_torch_list, + "torch_tensor": gen_root_link_state_torch_tensor, + }, + category="root_state", + ), + # --- Root Pose / Velocity --- + MethodBenchmark( + name="write_root_pose_to_sim", + method_name="write_root_pose_to_sim", + input_generators={ + "torch_list": gen_root_pose_torch_list, + "torch_tensor": gen_root_pose_torch_tensor, + }, + category="root_pose", + ), + MethodBenchmark( + name="write_root_link_pose_to_sim", + method_name="write_root_link_pose_to_sim", + input_generators={ + "torch_list": gen_root_link_pose_torch_list, + "torch_tensor": gen_root_link_pose_torch_tensor, + }, + category="root_pose", + ), + MethodBenchmark( + name="write_root_com_pose_to_sim", + method_name="write_root_com_pose_to_sim", + input_generators={ + "torch_list": gen_root_com_pose_torch_list, + "torch_tensor": gen_root_com_pose_torch_tensor, + }, + category="root_pose", + ), + MethodBenchmark( + name="write_root_velocity_to_sim", + method_name="write_root_velocity_to_sim", + input_generators={ + "torch_list": gen_root_velocity_torch_list, + "torch_tensor": gen_root_velocity_torch_tensor, + }, + category="root_velocity", + ), + MethodBenchmark( + name="write_root_link_velocity_to_sim", + method_name="write_root_link_velocity_to_sim", + input_generators={ + "torch_list": gen_root_link_velocity_torch_list, + "torch_tensor": gen_root_link_velocity_torch_tensor, + }, + category="root_velocity", + ), + MethodBenchmark( + name="write_root_com_velocity_to_sim", + method_name="write_root_com_velocity_to_sim", + input_generators={ + "torch_list": gen_root_com_velocity_torch_list, + "torch_tensor": gen_root_com_velocity_torch_tensor, + }, + category="root_velocity", + ), + # --- Body Properties --- + MethodBenchmark( + name="set_masses", + method_name="set_masses", + input_generators={ + "torch_list": gen_masses_torch_list, + "torch_tensor": gen_masses_torch_tensor, + }, + category="body_props", + ), + MethodBenchmark( + name="set_coms", + method_name="set_coms", + input_generators={ + "torch_list": gen_coms_torch_list, + "torch_tensor": gen_coms_torch_tensor, + }, + category="body_props", + ), + MethodBenchmark( + name="set_inertias", + method_name="set_inertias", + input_generators={ + "torch_list": gen_inertias_torch_list, + "torch_tensor": gen_inertias_torch_tensor, + }, + category="body_props", + ), + MethodBenchmark( + name="set_external_force_and_torque", + method_name="set_external_force_and_torque", + input_generators={ + "torch_list": gen_external_force_and_torque_torch_list, + "torch_tensor": gen_external_force_and_torque_torch_tensor, + }, + category="body_props", + ), +] + + +def run_benchmark(config: BenchmarkConfig): + """Run all benchmarks.""" + results = [] + + # Check if we should run all modes or specific ones + modes_to_run = [] + if isinstance(config.mode, str): + if config.mode == "all": + # Will be populated dynamically based on available generators + modes_to_run = None + else: + modes_to_run = [config.mode] + elif isinstance(config.mode, list): + modes_to_run = config.mode + + # Setup + rigid_object, mock_view, _ = create_test_rigid_object( + num_instances=config.num_instances, + num_bodies=config.num_bodies, + device=config.device, + ) + + print(f"Benchmarking RigidObject (PhysX) with {config.num_instances} instances, {config.num_bodies} bodies...") + print(f"Device: {config.device}") + print(f"Iterations: {config.num_iterations}, Warmup: {config.warmup_steps}") + print(f"Modes: {modes_to_run if modes_to_run else 'All available'}") + + print(f"\nBenchmarking {len(BENCHMARKS)} methods...") + for i, benchmark in enumerate(BENCHMARKS): + method = getattr(rigid_object, benchmark.method_name, None) + + # Determine which modes to run for this benchmark + available_modes = list(benchmark.input_generators.keys()) + current_modes = modes_to_run if modes_to_run is not None else available_modes + + # Filter modes that are available for this benchmark + current_modes = [m for m in current_modes if m in available_modes] + + for mode in current_modes: + generator = benchmark.input_generators[mode] + print(f"[{i + 1}/{len(BENCHMARKS)}] [{mode.upper()}] {benchmark.name}...", end=" ", flush=True) + + result = benchmark_method( + method=method, + method_name=benchmark.name, + generator=generator, + config=config, + dependencies=BENCHMARK_DEPENDENCIES, + ) + result.mode = mode + results.append(result) + + if result.skipped: + print(f"SKIPPED ({result.skip_reason})") + else: + print(f"{result.mean_time_us:.2f} ± {result.std_time_us:.2f} µs") + + return results + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Benchmark RigidObject methods (PhysX backend).") + parser.add_argument("--num_iterations", type=int, default=1000, help="Number of iterations") + parser.add_argument("--warmup_steps", type=int, default=10, help="Number of warmup steps") + parser.add_argument("--num_instances", type=int, default=4096, help="Number of instances") + parser.add_argument("--num_bodies", type=int, default=1, help="Number of bodies") + parser.add_argument("--device", type=str, default="cuda:0", help="Device") + parser.add_argument("--mode", type=str, default="all", help="Benchmark mode (all, torch_list, torch_tensor)") + parser.add_argument("--output", type=str, default=None, help="Output JSON filename") + parser.add_argument("--no_csv", action="store_true", help="Disable CSV output") + + args = parser.parse_args() + + config = BenchmarkConfig( + num_iterations=args.num_iterations, + warmup_steps=args.warmup_steps, + num_instances=args.num_instances, + num_bodies=args.num_bodies, + num_joints=0, + device=args.device, + mode=args.mode, + ) + + results = run_benchmark(config) + + hardware_info = get_hardware_info() + print_hardware_info(hardware_info) + print_results(results) + + if args.output: + json_filename = args.output + else: + json_filename = get_default_output_filename("rigid_object_benchmark") + + export_results_json(results, config, hardware_info, json_filename) + + if not args.no_csv: + csv_filename = json_filename.replace(".json", ".csv") + export_results_csv(results, csv_filename) diff --git a/source/isaaclab_physx/benchmark/assets/benchmark_rigid_object_collection.py b/source/isaaclab_physx/benchmark/assets/benchmark_rigid_object_collection.py new file mode 100644 index 00000000000..e62cc87579a --- /dev/null +++ b/source/isaaclab_physx/benchmark/assets/benchmark_rigid_object_collection.py @@ -0,0 +1,601 @@ +# 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 + +"""Micro-benchmarking framework for RigidObjectCollection class (PhysX backend). + +This module provides a benchmarking framework to measure the performance of setter and writer +methods in the RigidObjectCollection class. + +Usage: + python benchmark_rigid_object_collection.py [--num_iterations N] [--num_instances I] [--num_bodies B] + +Example: + python benchmark_rigid_object_collection.py --num_iterations 1000 --num_instances 4096 --num_bodies 4 +""" + +from __future__ import annotations + +import argparse +import sys +import warnings +from types import ModuleType +from unittest.mock import MagicMock + +import torch +import warp as wp + +# Initialize Warp +wp.init() + + +# ============================================================================= +# Mock Setup - Must happen BEFORE importing RigidObjectCollection +# ============================================================================= + + +class MockPhysicsSimView: + """Simple mock for the physics simulation view.""" + + def get_gravity(self): + return (0.0, 0.0, -9.81) + + +class MockSimulationManager: + """Simple mock for SimulationManager.""" + + @staticmethod + def get_physics_sim_view(): + return MockPhysicsSimView() + + +# Mock isaacsim.core.simulation_manager +mock_sim_manager_module = ModuleType("isaacsim.core.simulation_manager") +mock_sim_manager_module.SimulationManager = MockSimulationManager +sys.modules["isaacsim"] = ModuleType("isaacsim") +sys.modules["isaacsim.core"] = ModuleType("isaacsim.core") +sys.modules["isaacsim.core.simulation_manager"] = mock_sim_manager_module + +# Mock pxr (USD library) +sys.modules["pxr"] = MagicMock() +sys.modules["pxr.Usd"] = MagicMock() +sys.modules["pxr.UsdGeom"] = MagicMock() +sys.modules["pxr.UsdPhysics"] = MagicMock() + +# Mock omni module hierarchy +omni_mocks = [ + "omni", + "omni.kit", + "omni.kit.app", + "omni.kit.commands", + "omni.usd", + "omni.client", + "omni.timeline", + "omni.physx", + "omni.physx.scripts", + "omni.physx.scripts.utils", + "omni.physics", + "omni.physics.tensors", + "omni.physics.tensors.impl", + "omni.physics.tensors.impl.api", +] +for mod_name in omni_mocks: + mock = MagicMock() + mock.__name__ = mod_name + mock.__path__ = [] + mock.__package__ = mod_name + sys.modules[mod_name] = mock + +# Mock carb +mock_carb = MagicMock() +mock_carb.settings.get_settings.return_value.get.return_value = "/mock/path" +sys.modules["carb"] = mock_carb + +# Mock isaaclab.sim module hierarchy +sim_mock = MagicMock() +sim_mock.find_first_matching_prim = MagicMock() +sim_mock.get_all_matching_child_prims = MagicMock(return_value=[]) +sys.modules["isaaclab.sim"] = sim_mock +sys.modules["isaaclab.sim.utils"] = MagicMock() +sys.modules["isaaclab.sim.utils.stage"] = MagicMock() +sys.modules["isaaclab.sim.converters"] = MagicMock() + +# Mock WrenchComposer - import from mock_interfaces and patch into the module +from isaaclab.test.mock_interfaces.utils import MockWrenchComposer + +mock_wrench_module = ModuleType("isaaclab.utils.wrench_composer") +mock_wrench_module.WrenchComposer = MockWrenchComposer +sys.modules["isaaclab.utils.wrench_composer"] = mock_wrench_module + + +# Mock base classes +class BaseRigidObjectCollection: + """Mock base class.""" + + @property + def device(self) -> str: + return self._device + + +class BaseRigidObjectCollectionData: + """Mock base class.""" + + def __init__(self, root_view, num_bodies: int, device: str): + self.device = device + + +mock_base_collection = ModuleType("isaaclab.assets.rigid_object_collection.base_rigid_object_collection") +mock_base_collection.BaseRigidObjectCollection = BaseRigidObjectCollection +sys.modules["isaaclab.assets.rigid_object_collection.base_rigid_object_collection"] = mock_base_collection + +mock_base_collection_data = ModuleType("isaaclab.assets.rigid_object_collection.base_rigid_object_collection_data") +mock_base_collection_data.BaseRigidObjectCollectionData = BaseRigidObjectCollectionData +sys.modules["isaaclab.assets.rigid_object_collection.base_rigid_object_collection_data"] = mock_base_collection_data + +# Now import via importlib +import importlib.util +from pathlib import Path + +benchmark_dir = Path(__file__).resolve().parent + +# Load RigidObjectCollectionData +data_path = ( + benchmark_dir.parents[1] + / "isaaclab_physx" + / "assets" + / "rigid_object_collection" + / "rigid_object_collection_data.py" +) +spec = importlib.util.spec_from_file_location( + "isaaclab_physx.assets.rigid_object_collection.rigid_object_collection_data", data_path +) +data_module = importlib.util.module_from_spec(spec) +sys.modules["isaaclab_physx.assets.rigid_object_collection.rigid_object_collection_data"] = data_module +spec.loader.exec_module(data_module) +RigidObjectCollectionData = data_module.RigidObjectCollectionData + +# Load RigidObjectCollection +collection_path = ( + benchmark_dir.parents[1] / "isaaclab_physx" / "assets" / "rigid_object_collection" / "rigid_object_collection.py" +) +spec = importlib.util.spec_from_file_location( + "isaaclab_physx.assets.rigid_object_collection.rigid_object_collection", collection_path +) +collection_module = importlib.util.module_from_spec(spec) +sys.modules["isaaclab_physx.assets.rigid_object_collection.rigid_object_collection"] = collection_module +spec.loader.exec_module(collection_module) +RigidObjectCollection = collection_module.RigidObjectCollection + +# Import shared utilities +# Import mock view +from isaaclab_physx.test.mock_interfaces.views import MockRigidBodyView + +from isaaclab.test.benchmark import ( + BenchmarkConfig, + MethodBenchmark, + benchmark_method, + export_results_csv, + export_results_json, + get_default_output_filename, + get_hardware_info, + make_tensor_body_ids, + make_tensor_env_ids, + print_hardware_info, + print_results, +) + +# Suppress warnings +warnings.filterwarnings("ignore", category=DeprecationWarning) +warnings.filterwarnings("ignore", category=UserWarning) + +# Suppress isaaclab logging (deprecation warnings) +import logging + +logging.getLogger("isaaclab_physx").setLevel(logging.ERROR) +logging.getLogger("isaaclab").setLevel(logging.ERROR) + + +# Simple config class +class RigidObjectCollectionCfg: + def __init__(self, prim_path: str = "/World/Collection"): + self.prim_path = prim_path + + +def create_test_collection( + num_instances: int = 2, + num_bodies: int = 4, + device: str = "cuda:0", +) -> tuple[RigidObjectCollection, MockRigidBodyView]: + """Create a test RigidObjectCollection instance with mocked dependencies.""" + object_names = [f"object_{i}" for i in range(num_bodies)] + + collection = object.__new__(RigidObjectCollection) + + collection.cfg = RigidObjectCollectionCfg(prim_path="/World/Collection") + + # Create PhysX mock view (total count = num_instances * num_bodies) + total_count = num_instances * num_bodies + mock_view = MockRigidBodyView( + count=total_count, + device=device, + ) + mock_view.set_random_mock_data() + + # Set up attributes required before _create_buffers + object.__setattr__(collection, "_root_view", mock_view) + object.__setattr__(collection, "_device", device) + object.__setattr__(collection, "_body_names_list", object_names) + + # Create RigidObjectCollectionData instance + data = RigidObjectCollectionData(mock_view, num_bodies, device) + data.object_names = object_names + object.__setattr__(collection, "_data", data) + + # Call _create_buffers to set up all internal buffers and wrench composers + collection._create_buffers() + + return collection, mock_view + + +# ============================================================================= +# Input Generators +# ============================================================================= + + +def gen_body_state_torch_list(config: BenchmarkConfig) -> dict: + return { + "body_states": torch.rand( + config.num_instances, config.num_bodies, 13, device=config.device, dtype=torch.float32 + ), + "env_ids": list(range(config.num_instances)), + "body_ids": list(range(config.num_bodies)), + } + + +def gen_body_state_torch_tensor(config: BenchmarkConfig) -> dict: + return { + "body_states": torch.rand( + config.num_instances, config.num_bodies, 13, device=config.device, dtype=torch.float32 + ), + "env_ids": make_tensor_env_ids(config.num_instances, config.device), + "body_ids": make_tensor_body_ids(config.num_bodies, config.device), + } + + +def gen_body_pose_torch_list(config: BenchmarkConfig) -> dict: + return { + "body_poses": torch.rand(config.num_instances, config.num_bodies, 7, device=config.device, dtype=torch.float32), + "env_ids": list(range(config.num_instances)), + "body_ids": list(range(config.num_bodies)), + } + + +def gen_body_pose_torch_tensor(config: BenchmarkConfig) -> dict: + return { + "body_poses": torch.rand(config.num_instances, config.num_bodies, 7, device=config.device, dtype=torch.float32), + "env_ids": make_tensor_env_ids(config.num_instances, config.device), + "body_ids": make_tensor_body_ids(config.num_bodies, config.device), + } + + +def gen_body_velocity_torch_list(config: BenchmarkConfig) -> dict: + return { + "body_velocities": torch.rand( + config.num_instances, config.num_bodies, 6, device=config.device, dtype=torch.float32 + ), + "env_ids": list(range(config.num_instances)), + "body_ids": list(range(config.num_bodies)), + } + + +def gen_body_velocity_torch_tensor(config: BenchmarkConfig) -> dict: + return { + "body_velocities": torch.rand( + config.num_instances, config.num_bodies, 6, device=config.device, dtype=torch.float32 + ), + "env_ids": make_tensor_env_ids(config.num_instances, config.device), + "body_ids": make_tensor_body_ids(config.num_bodies, config.device), + } + + +# --- External Force and Torque --- +def gen_external_force_and_torque_torch_list(config: BenchmarkConfig) -> dict: + return { + "forces": torch.rand(config.num_instances, config.num_bodies, 3, device=config.device, dtype=torch.float32), + "torques": torch.rand(config.num_instances, config.num_bodies, 3, device=config.device, dtype=torch.float32), + "env_ids": list(range(config.num_instances)), + } + + +def gen_external_force_and_torque_torch_tensor(config: BenchmarkConfig) -> dict: + return { + "forces": torch.rand(config.num_instances, config.num_bodies, 3, device=config.device, dtype=torch.float32), + "torques": torch.rand(config.num_instances, config.num_bodies, 3, device=config.device, dtype=torch.float32), + "env_ids": make_tensor_env_ids(config.num_instances, config.device), + } + + +# --- Masses --- +def gen_masses_torch_list(config: BenchmarkConfig) -> dict: + return { + "masses": torch.rand(config.num_instances, config.num_bodies, device=config.device, dtype=torch.float32), + "env_ids": list(range(config.num_instances)), + "body_ids": list(range(config.num_bodies)), + } + + +def gen_masses_torch_tensor(config: BenchmarkConfig) -> dict: + return { + "masses": torch.rand(config.num_instances, config.num_bodies, device=config.device, dtype=torch.float32), + "env_ids": make_tensor_env_ids(config.num_instances, config.device), + "body_ids": make_tensor_body_ids(config.num_bodies, config.device), + } + + +# --- CoMs --- +def gen_coms_torch_list(config: BenchmarkConfig) -> dict: + return { + "coms": torch.rand(config.num_instances, config.num_bodies, 3, device=config.device, dtype=torch.float32), + "env_ids": list(range(config.num_instances)), + "body_ids": list(range(config.num_bodies)), + } + + +def gen_coms_torch_tensor(config: BenchmarkConfig) -> dict: + return { + "coms": torch.rand(config.num_instances, config.num_bodies, 3, device=config.device, dtype=torch.float32), + "env_ids": make_tensor_env_ids(config.num_instances, config.device), + "body_ids": make_tensor_body_ids(config.num_bodies, config.device), + } + + +# --- Inertias --- +def gen_inertias_torch_list(config: BenchmarkConfig) -> dict: + return { + "inertias": torch.rand( + config.num_instances, config.num_bodies, 3, 3, device=config.device, dtype=torch.float32 + ), + "env_ids": list(range(config.num_instances)), + "body_ids": list(range(config.num_bodies)), + } + + +def gen_inertias_torch_tensor(config: BenchmarkConfig) -> dict: + return { + "inertias": torch.rand( + config.num_instances, config.num_bodies, 3, 3, device=config.device, dtype=torch.float32 + ), + "env_ids": make_tensor_env_ids(config.num_instances, config.device), + "body_ids": make_tensor_body_ids(config.num_bodies, config.device), + } + + +# ============================================================================= +# Benchmarks +# ============================================================================= + +BENCHMARK_DEPENDENCIES = {} + +BENCHMARKS = [ + # --- Body State --- + MethodBenchmark( + name="write_body_state_to_sim", + method_name="write_body_state_to_sim", + input_generators={ + "torch_list": gen_body_state_torch_list, + "torch_tensor": gen_body_state_torch_tensor, + }, + category="body_state", + ), + MethodBenchmark( + name="write_body_link_state_to_sim", + method_name="write_body_link_state_to_sim", + input_generators={ + "torch_list": gen_body_state_torch_list, + "torch_tensor": gen_body_state_torch_tensor, + }, + category="body_state", + ), + MethodBenchmark( + name="write_body_com_state_to_sim", + method_name="write_body_com_state_to_sim", + input_generators={ + "torch_list": gen_body_state_torch_list, + "torch_tensor": gen_body_state_torch_tensor, + }, + category="body_state", + ), + # --- Body Pose --- + MethodBenchmark( + name="write_body_pose_to_sim", + method_name="write_body_pose_to_sim", + input_generators={ + "torch_list": gen_body_pose_torch_list, + "torch_tensor": gen_body_pose_torch_tensor, + }, + category="body_pose", + ), + MethodBenchmark( + name="write_body_link_pose_to_sim", + method_name="write_body_link_pose_to_sim", + input_generators={ + "torch_list": gen_body_pose_torch_list, + "torch_tensor": gen_body_pose_torch_tensor, + }, + category="body_pose", + ), + MethodBenchmark( + name="write_body_com_pose_to_sim", + method_name="write_body_com_pose_to_sim", + input_generators={ + "torch_list": gen_body_pose_torch_list, + "torch_tensor": gen_body_pose_torch_tensor, + }, + category="body_pose", + ), + # --- Body Velocity --- + MethodBenchmark( + name="write_body_velocity_to_sim", + method_name="write_body_velocity_to_sim", + input_generators={ + "torch_list": gen_body_velocity_torch_list, + "torch_tensor": gen_body_velocity_torch_tensor, + }, + category="body_velocity", + ), + MethodBenchmark( + name="write_body_link_velocity_to_sim", + method_name="write_body_link_velocity_to_sim", + input_generators={ + "torch_list": gen_body_velocity_torch_list, + "torch_tensor": gen_body_velocity_torch_tensor, + }, + category="body_velocity", + ), + MethodBenchmark( + name="write_body_com_velocity_to_sim", + method_name="write_body_com_velocity_to_sim", + input_generators={ + "torch_list": gen_body_velocity_torch_list, + "torch_tensor": gen_body_velocity_torch_tensor, + }, + category="body_velocity", + ), + # --- External Force and Torque --- + MethodBenchmark( + name="set_external_force_and_torque", + method_name="set_external_force_and_torque", + input_generators={ + "torch_list": gen_external_force_and_torque_torch_list, + "torch_tensor": gen_external_force_and_torque_torch_tensor, + }, + category="external_wrench", + ), + # --- Body Properties --- + MethodBenchmark( + name="set_masses", + method_name="set_masses", + input_generators={ + "torch_list": gen_masses_torch_list, + "torch_tensor": gen_masses_torch_tensor, + }, + category="body_props", + ), + MethodBenchmark( + name="set_coms", + method_name="set_coms", + input_generators={ + "torch_list": gen_coms_torch_list, + "torch_tensor": gen_coms_torch_tensor, + }, + category="body_props", + ), + MethodBenchmark( + name="set_inertias", + method_name="set_inertias", + input_generators={ + "torch_list": gen_inertias_torch_list, + "torch_tensor": gen_inertias_torch_tensor, + }, + category="body_props", + ), +] + + +def run_benchmark(config: BenchmarkConfig): + """Run all benchmarks.""" + results = [] + + modes_to_run = [] + if isinstance(config.mode, str): + if config.mode == "all": + modes_to_run = None + else: + modes_to_run = [config.mode] + elif isinstance(config.mode, list): + modes_to_run = config.mode + + collection, mock_view = create_test_collection( + num_instances=config.num_instances, + num_bodies=config.num_bodies, + device=config.device, + ) + + print( + f"Benchmarking RigidObjectCollection (PhysX) with {config.num_instances} instances, " + f"{config.num_bodies} bodies..." + ) + print(f"Device: {config.device}") + print(f"Iterations: {config.num_iterations}, Warmup: {config.warmup_steps}") + print(f"Modes: {modes_to_run if modes_to_run else 'All available'}") + + print(f"\nBenchmarking {len(BENCHMARKS)} methods...") + for i, benchmark in enumerate(BENCHMARKS): + method = getattr(collection, benchmark.method_name, None) + + available_modes = list(benchmark.input_generators.keys()) + current_modes = modes_to_run if modes_to_run is not None else available_modes + current_modes = [m for m in current_modes if m in available_modes] + + for mode in current_modes: + generator = benchmark.input_generators[mode] + print(f"[{i + 1}/{len(BENCHMARKS)}] [{mode.upper()}] {benchmark.name}...", end=" ", flush=True) + + result = benchmark_method( + method=method, + method_name=benchmark.name, + generator=generator, + config=config, + dependencies=BENCHMARK_DEPENDENCIES, + ) + result.mode = mode + results.append(result) + + if result.skipped: + print(f"SKIPPED ({result.skip_reason})") + else: + print(f"{result.mean_time_us:.2f} ± {result.std_time_us:.2f} µs") + + return results + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Benchmark RigidObjectCollection methods (PhysX backend).") + parser.add_argument("--num_iterations", type=int, default=1000, help="Number of iterations") + parser.add_argument("--warmup_steps", type=int, default=10, help="Number of warmup steps") + parser.add_argument("--num_instances", type=int, default=4096, help="Number of instances") + parser.add_argument("--num_bodies", type=int, default=4, help="Number of bodies per instance") + parser.add_argument("--device", type=str, default="cuda:0", help="Device") + parser.add_argument("--mode", type=str, default="all", help="Benchmark mode (all, torch_list, torch_tensor)") + parser.add_argument("--output", type=str, default=None, help="Output JSON filename") + parser.add_argument("--no_csv", action="store_true", help="Disable CSV output") + + args = parser.parse_args() + + config = BenchmarkConfig( + num_iterations=args.num_iterations, + warmup_steps=args.warmup_steps, + num_instances=args.num_instances, + num_bodies=args.num_bodies, + num_joints=0, + device=args.device, + mode=args.mode, + ) + + results = run_benchmark(config) + + hardware_info = get_hardware_info() + print_hardware_info(hardware_info) + print_results(results) + + if args.output: + json_filename = args.output + else: + json_filename = get_default_output_filename("rigid_object_collection_benchmark") + + export_results_json(results, config, hardware_info, json_filename) + + if not args.no_csv: + csv_filename = json_filename.replace(".json", ".csv") + export_results_csv(results, csv_filename) diff --git a/source/isaaclab_physx/benchmark/assets/benchmark_rigid_object_collection_data.py b/source/isaaclab_physx/benchmark/assets/benchmark_rigid_object_collection_data.py new file mode 100644 index 00000000000..f715b681b39 --- /dev/null +++ b/source/isaaclab_physx/benchmark/assets/benchmark_rigid_object_collection_data.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 + +"""Micro-benchmarking framework for RigidObjectCollectionData class. + +This module provides a benchmarking framework to measure the performance of all functions +in the RigidObjectCollectionData class. Each function is run multiple times with randomized mock data, +and timing statistics (mean and standard deviation) are reported. + +Usage: + python benchmark_rigid_object_collection_data.py [--num_iterations N] [--warmup_steps W] + [--num_instances I] [--num_bodies B] + +Example: + python benchmark_rigid_object_collection_data.py --num_iterations 10000 --warmup_steps 10 +""" + +from __future__ import annotations + +import argparse +import sys +import warnings +from types import ModuleType +from unittest.mock import MagicMock + +import torch +import warp as wp + +# Initialize Warp first +wp.init() + + +# ============================================================================= +# Mock Setup - Must happen BEFORE importing RigidObjectCollectionData +# ============================================================================= + + +# Mock BaseRigidObjectCollectionData - this is just an abstract class +class BaseRigidObjectCollectionData: + """Mock base class to avoid importing isaaclab.assets (which has many dependencies).""" + + def __init__(self, root_view, num_bodies: int, device: str): + self.device = device + + +# Create mock module for isaaclab.assets.rigid_object_collection.base_rigid_object_collection_data +mock_base_module = ModuleType("isaaclab.assets.rigid_object_collection.base_rigid_object_collection_data") +mock_base_module.BaseRigidObjectCollectionData = BaseRigidObjectCollectionData +sys.modules["isaaclab.assets.rigid_object_collection.base_rigid_object_collection_data"] = mock_base_module + +# Mock pxr (USD library - not available in headless docker, used by isaaclab.utils.mesh) +sys.modules["pxr"] = MagicMock() +sys.modules["pxr.Usd"] = MagicMock() +sys.modules["pxr.UsdGeom"] = MagicMock() + + +class MockPhysicsSimView: + """Simple mock for the physics simulation view.""" + + def get_gravity(self): + """Return gravity as a tuple of 3 floats.""" + return (0.0, 0.0, -9.81) + + +class MockSimulationManager: + """Simple mock for SimulationManager.""" + + @staticmethod + def get_physics_sim_view(): + return MockPhysicsSimView() + + +# Mock isaacsim.core.simulation_manager +mock_sim_manager_module = ModuleType("isaacsim.core.simulation_manager") +mock_sim_manager_module.SimulationManager = MockSimulationManager +sys.modules["isaacsim"] = ModuleType("isaacsim") +sys.modules["isaacsim.core"] = ModuleType("isaacsim.core") +sys.modules["isaacsim.core.simulation_manager"] = mock_sim_manager_module + +# Now we can directly import RigidObjectCollectionData +import importlib.util +from pathlib import Path + +benchmark_dir = Path(__file__).resolve().parent +data_path = ( + benchmark_dir.parents[1] + / "isaaclab_physx" + / "assets" + / "rigid_object_collection" + / "rigid_object_collection_data.py" +) + +spec = importlib.util.spec_from_file_location( + "isaaclab_physx.assets.rigid_object_collection.rigid_object_collection_data", data_path +) +data_module = importlib.util.module_from_spec(spec) +sys.modules["isaaclab_physx.assets.rigid_object_collection.rigid_object_collection_data"] = data_module +spec.loader.exec_module(data_module) +RigidObjectCollectionData = data_module.RigidObjectCollectionData + +# Import shared utilities from common module +# Import mock classes from PhysX test utilities +from isaaclab_physx.test.mock_interfaces.views import MockRigidBodyView + +from isaaclab.test.benchmark import ( + BenchmarkConfig, + BenchmarkResult, + MethodBenchmark, + benchmark_method, + export_results_csv, + export_results_json, + get_default_output_filename, + get_hardware_info, + print_hardware_info, + print_results, +) + +# Suppress deprecation warnings during benchmarking +warnings.filterwarnings("ignore", category=DeprecationWarning) +warnings.filterwarnings("ignore", category=UserWarning) + + +# ============================================================================= +# Skip Lists +# ============================================================================= + +# List of deprecated properties (for backward compatibility) - skip these +DEPRECATED_PROPERTIES = { + "default_root_state", + "root_pose_w", + "root_pos_w", + "root_quat_w", + "root_vel_w", + "root_lin_vel_w", + "root_ang_vel_w", + "root_lin_vel_b", + "root_ang_vel_b", + "body_pose_w", + "body_pos_w", + "body_quat_w", + "body_vel_w", + "body_lin_vel_w", + "body_ang_vel_w", + "body_acc_w", + "body_lin_acc_w", + "body_ang_acc_w", + "com_pos_b", + "com_quat_b", + # Also skip the combined state properties marked as deprecated + "root_state_w", + "object_link_state_w", + "object_com_state_w", +} + +# List of properties that raise NotImplementedError - skip these +NOT_IMPLEMENTED_PROPERTIES = {} + +# Removed default_* properties that raise RuntimeError +REMOVED_PROPERTIES = { + "default_inertia", + "default_mass", +} + +# Private/internal properties and methods to skip +INTERNAL_PROPERTIES = { + "_create_simulation_bindings", + "_create_buffers", + "update", + "is_primed", + "device", + "object_names", + "GRAVITY_VEC_W", + "GRAVITY_VEC_W_TORCH", + "FORWARD_VEC_B", + "FORWARD_VEC_B_TORCH", + "ALL_ENV_MASK", + "ENV_MASK", + "ALL_OBJECT_MASK", + "OBJECT_MASK", + "num_bodies", + "num_instances", +} + +# Dependency mapping for properties +PROPERTY_DEPENDENCIES = { + "object_link_lin_vel_w": ["object_link_vel_w"], + "object_link_ang_vel_w": ["object_link_vel_w"], + "object_link_lin_vel_b": ["object_link_vel_b"], + "object_link_ang_vel_b": ["object_link_vel_b"], + "object_com_pos_w": ["object_com_pose_w"], + "object_com_quat_w": ["object_com_pose_w"], + "object_com_lin_vel_b": ["object_com_vel_b"], + "object_com_ang_vel_b": ["object_com_vel_b"], + "object_com_lin_vel_w": ["object_com_vel_w"], + "object_com_ang_vel_w": ["object_com_vel_w"], + "object_link_pos_w": ["object_link_pose_w"], + "object_link_quat_w": ["object_link_pose_w"], + "object_com_lin_acc_w": ["object_com_acc_w"], + "object_com_ang_acc_w": ["object_com_acc_w"], + "object_com_quat_b": ["object_com_pose_b"], +} + + +# ============================================================================= +# Benchmark Functions +# ============================================================================= + + +def get_benchmarkable_properties(data: RigidObjectCollectionData) -> list[str]: + """Get list of properties that can be benchmarked.""" + all_properties = [] + + for name in dir(data): + if name.startswith("_"): + continue + if name in DEPRECATED_PROPERTIES: + continue + if name in NOT_IMPLEMENTED_PROPERTIES: + continue + if name in REMOVED_PROPERTIES: + continue + if name in INTERNAL_PROPERTIES: + continue + + try: + attr = getattr(type(data), name, None) + if isinstance(attr, property): + all_properties.append(name) + except Exception: + pass + + return sorted(all_properties) + + +def setup_mock_environment(config: BenchmarkConfig) -> MockRigidBodyView: + """Set up the mock environment for benchmarking.""" + # For collection, total count = num_instances * num_bodies + total_count = config.num_instances * config.num_bodies + mock_view = MockRigidBodyView( + count=total_count, + device=config.device, + ) + return mock_view + + +def run_benchmarks(config: BenchmarkConfig) -> list[BenchmarkResult]: + """Run all benchmarks for RigidObjectCollectionData.""" + results = [] + + # Setup mock environment + mock_view = setup_mock_environment(config) + mock_view.set_random_mock_data() + + # Create RigidObjectCollectionData instance + data = RigidObjectCollectionData(mock_view, config.num_bodies, config.device) + + # Get list of properties to benchmark + properties = get_benchmarkable_properties(data) + + total_count = config.num_instances * config.num_bodies + + # Generator that updates mock data and invalidates timestamp + def gen_mock_data(cfg: BenchmarkConfig) -> dict: + mock_view.set_mock_transforms(torch.randn(total_count, 7, device=cfg.device)) + mock_view.set_mock_velocities(torch.randn(total_count, 6, device=cfg.device)) + mock_view.set_mock_accelerations(torch.randn(total_count, 6, device=cfg.device)) + mock_view.set_mock_coms(torch.randn(total_count, 7, device=cfg.device)) + data._sim_timestamp += 1.0 + return {} + + # Create benchmarks dynamically + benchmarks = [] + for prop_name in properties: + benchmarks.append( + MethodBenchmark( + name=prop_name, + method_name=prop_name, + input_generators={"default": gen_mock_data}, + category="property", + ) + ) + + print(f"\nBenchmarking {len(benchmarks)} properties...") + print(f"Config: {config.num_iterations} iterations, {config.warmup_steps} warmup steps") + print(f" {config.num_instances} instances, {config.num_bodies} bodies") + print("-" * 80) + + for i, benchmark in enumerate(benchmarks): + + def prop_accessor(prop=benchmark.method_name, **kwargs): + return getattr(data, prop) + + print(f"[{i + 1}/{len(benchmarks)}] [DEFAULT] {benchmark.name}...", end=" ", flush=True) + + result = benchmark_method( + method=prop_accessor, + method_name=benchmark.name, + generator=gen_mock_data, + config=config, + dependencies=PROPERTY_DEPENDENCIES, + ) + result.mode = "default" + results.append(result) + + if result.skipped: + print(f"SKIPPED ({result.skip_reason})") + else: + print(f"{result.mean_time_us:.2f} ± {result.std_time_us:.2f} µs") + + return results + + +def main(): + """Main entry point for the benchmarking script.""" + parser = argparse.ArgumentParser( + description="Micro-benchmarking framework for RigidObjectCollectionData class.", + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + ) + parser.add_argument("--num_iterations", type=int, default=1000, help="Number of iterations") + parser.add_argument("--warmup_steps", type=int, default=10, help="Number of warmup steps") + parser.add_argument("--num_instances", type=int, default=4096, help="Number of instances") + parser.add_argument("--num_bodies", type=int, default=4, help="Number of bodies per instance") + parser.add_argument("--device", type=str, default="cuda:0", help="Device") + parser.add_argument("--output", type=str, default=None, help="Output JSON filename") + parser.add_argument("--no_csv", action="store_true", help="Disable CSV output") + + args = parser.parse_args() + + config = BenchmarkConfig( + num_iterations=args.num_iterations, + warmup_steps=args.warmup_steps, + num_instances=args.num_instances, + num_bodies=args.num_bodies, + num_joints=0, + device=args.device, + ) + + results = run_benchmarks(config) + + hardware_info = get_hardware_info() + print_hardware_info(hardware_info) + print_results(results, include_mode=False) + + if args.output: + json_filename = args.output + else: + json_filename = get_default_output_filename("rigid_object_collection_data_benchmark") + + export_results_json(results, config, hardware_info, json_filename) + + if not args.no_csv: + csv_filename = json_filename.replace(".json", ".csv") + export_results_csv(results, csv_filename) + + +if __name__ == "__main__": + main() diff --git a/source/isaaclab_physx/benchmark/assets/benchmark_rigid_object_data.py b/source/isaaclab_physx/benchmark/assets/benchmark_rigid_object_data.py new file mode 100644 index 00000000000..99a2d85b93f --- /dev/null +++ b/source/isaaclab_physx/benchmark/assets/benchmark_rigid_object_data.py @@ -0,0 +1,357 @@ +# 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 + +"""Micro-benchmarking framework for RigidObjectData class. + +This module provides a benchmarking framework to measure the performance of all functions +in the RigidObjectData class. Each function is run multiple times with randomized mock data, +and timing statistics (mean and standard deviation) are reported. + +Usage: + python benchmark_rigid_object_data.py [--num_iterations N] [--warmup_steps W] [--num_instances I] + +Example: + python benchmark_rigid_object_data.py --num_iterations 10000 --warmup_steps 10 +""" + +from __future__ import annotations + +import argparse +import sys +import warnings +from types import ModuleType +from unittest.mock import MagicMock + +import torch +import warp as wp + +# Initialize Warp first +wp.init() + + +# ============================================================================= +# Mock Setup - Must happen BEFORE importing RigidObjectData +# ============================================================================= + + +# Mock BaseRigidObjectData - this is just an abstract class, so we provide an empty one +# to avoid importing the whole isaaclab.assets package +class BaseRigidObjectData: + """Mock base class to avoid importing isaaclab.assets (which has many dependencies).""" + + def __init__(self, root_view, device: str): + self.device = device + + +# Create mock module for isaaclab.assets.rigid_object.base_rigid_object_data +mock_base_module = ModuleType("isaaclab.assets.rigid_object.base_rigid_object_data") +mock_base_module.BaseRigidObjectData = BaseRigidObjectData +sys.modules["isaaclab.assets.rigid_object.base_rigid_object_data"] = mock_base_module + +# Mock pxr (USD library - not available in headless docker, used by isaaclab.utils.mesh) +sys.modules["pxr"] = MagicMock() +sys.modules["pxr.Usd"] = MagicMock() +sys.modules["pxr.UsdGeom"] = MagicMock() + + +class MockPhysicsSimView: + """Simple mock for the physics simulation view.""" + + def get_gravity(self): + """Return gravity as a tuple of 3 floats.""" + return (0.0, 0.0, -9.81) + + +class MockSimulationManager: + """Simple mock for SimulationManager.""" + + @staticmethod + def get_physics_sim_view(): + return MockPhysicsSimView() + + +# Mock isaacsim.core.simulation_manager +mock_sim_manager_module = ModuleType("isaacsim.core.simulation_manager") +mock_sim_manager_module.SimulationManager = MockSimulationManager +sys.modules["isaacsim"] = ModuleType("isaacsim") +sys.modules["isaacsim.core"] = ModuleType("isaacsim.core") +sys.modules["isaacsim.core.simulation_manager"] = mock_sim_manager_module + +# Now we can directly import RigidObjectData +import importlib.util +from pathlib import Path + +benchmark_dir = Path(__file__).resolve().parent +rigid_object_data_path = ( + benchmark_dir.parents[1] / "isaaclab_physx" / "assets" / "rigid_object" / "rigid_object_data.py" +) + +spec = importlib.util.spec_from_file_location( + "isaaclab_physx.assets.rigid_object.rigid_object_data", rigid_object_data_path +) +rigid_object_data_module = importlib.util.module_from_spec(spec) +sys.modules["isaaclab_physx.assets.rigid_object.rigid_object_data"] = rigid_object_data_module +spec.loader.exec_module(rigid_object_data_module) +RigidObjectData = rigid_object_data_module.RigidObjectData + +# Import shared utilities from common module +# Import mock classes from PhysX test utilities +from isaaclab_physx.test.mock_interfaces.views import MockRigidBodyView + +from isaaclab.test.benchmark import ( + BenchmarkConfig, + BenchmarkResult, + MethodBenchmark, + benchmark_method, + export_results_csv, + export_results_json, + get_default_output_filename, + get_hardware_info, + print_hardware_info, + print_results, +) + +# Suppress deprecation warnings during benchmarking +warnings.filterwarnings("ignore", category=DeprecationWarning) +warnings.filterwarnings("ignore", category=UserWarning) + + +# ============================================================================= +# Skip Lists +# ============================================================================= + +# List of deprecated properties (for backward compatibility) - skip these +DEPRECATED_PROPERTIES = { + "default_root_state", + "root_pose_w", + "root_pos_w", + "root_quat_w", + "root_vel_w", + "root_lin_vel_w", + "root_ang_vel_w", + "root_lin_vel_b", + "root_ang_vel_b", + "body_pose_w", + "body_pos_w", + "body_quat_w", + "body_vel_w", + "body_lin_vel_w", + "body_ang_vel_w", + "body_acc_w", + "body_lin_acc_w", + "body_ang_acc_w", + "com_pos_b", + "com_quat_b", + # Also skip the combined state properties marked as deprecated + "root_state_w", + "root_link_state_w", + "root_com_state_w", + "body_state_w", + "body_link_state_w", + "body_com_state_w", +} + +# List of properties that raise NotImplementedError - skip these +NOT_IMPLEMENTED_PROPERTIES = {} + +# Removed default_* properties that raise RuntimeError +REMOVED_PROPERTIES = { + "default_inertia", + "default_mass", +} + +# Private/internal properties and methods to skip +INTERNAL_PROPERTIES = { + "_create_simulation_bindings", + "_create_buffers", + "update", + "is_primed", + "device", + "body_names", + "GRAVITY_VEC_W", + "GRAVITY_VEC_W_TORCH", + "FORWARD_VEC_B", + "FORWARD_VEC_B_TORCH", + "ALL_ENV_MASK", + "ENV_MASK", +} + +# Dependency mapping for properties +PROPERTY_DEPENDENCIES = { + "root_link_lin_vel_w": ["root_link_vel_w"], + "root_link_ang_vel_w": ["root_link_vel_w"], + "root_link_lin_vel_b": ["root_link_vel_b"], + "root_link_ang_vel_b": ["root_link_vel_b"], + "root_com_pos_w": ["root_com_pose_w"], + "root_com_quat_w": ["root_com_pose_w"], + "root_com_lin_vel_b": ["root_com_vel_b"], + "root_com_ang_vel_b": ["root_com_vel_b"], + "root_com_lin_vel_w": ["root_com_vel_w"], + "root_com_ang_vel_w": ["root_com_vel_w"], + "root_link_pos_w": ["root_link_pose_w"], + "root_link_quat_w": ["root_link_pose_w"], + "body_link_lin_vel_w": ["body_link_vel_w"], + "body_link_ang_vel_w": ["body_link_vel_w"], + "body_link_pos_w": ["body_link_pose_w"], + "body_link_quat_w": ["body_link_pose_w"], + "body_com_pos_w": ["body_com_pose_w"], + "body_com_quat_w": ["body_com_pose_w"], + "body_com_lin_vel_w": ["body_com_vel_w"], + "body_com_ang_vel_w": ["body_com_vel_w"], + "body_com_lin_acc_w": ["body_com_acc_w"], + "body_com_ang_acc_w": ["body_com_acc_w"], + "body_com_quat_b": ["body_com_pose_b"], +} + + +# ============================================================================= +# Benchmark Functions +# ============================================================================= + + +def get_benchmarkable_properties(rigid_object_data: RigidObjectData) -> list[str]: + """Get list of properties that can be benchmarked.""" + all_properties = [] + + for name in dir(rigid_object_data): + if name.startswith("_"): + continue + if name in DEPRECATED_PROPERTIES: + continue + if name in NOT_IMPLEMENTED_PROPERTIES: + continue + if name in REMOVED_PROPERTIES: + continue + if name in INTERNAL_PROPERTIES: + continue + + try: + attr = getattr(type(rigid_object_data), name, None) + if isinstance(attr, property): + all_properties.append(name) + except Exception: + pass + + return sorted(all_properties) + + +def setup_mock_environment(config: BenchmarkConfig) -> MockRigidBodyView: + """Set up the mock environment for benchmarking.""" + mock_view = MockRigidBodyView( + count=config.num_instances, + device=config.device, + ) + return mock_view + + +def run_benchmarks(config: BenchmarkConfig) -> list[BenchmarkResult]: + """Run all benchmarks for RigidObjectData.""" + results = [] + + # Setup mock environment + mock_view = setup_mock_environment(config) + mock_view.set_random_mock_data() + + # Create RigidObjectData instance + rigid_object_data = RigidObjectData(mock_view, config.device) + + # Get list of properties to benchmark + properties = get_benchmarkable_properties(rigid_object_data) + + # Generator that updates mock data and invalidates timestamp + def gen_mock_data(cfg: BenchmarkConfig) -> dict: + mock_view.set_mock_transforms(torch.randn(cfg.num_instances, 7, device=cfg.device)) + mock_view.set_mock_velocities(torch.randn(cfg.num_instances, 6, device=cfg.device)) + mock_view.set_mock_accelerations(torch.randn(cfg.num_instances, 6, device=cfg.device)) + mock_view.set_mock_coms(torch.randn(cfg.num_instances, 7, device=cfg.device)) + rigid_object_data._sim_timestamp += 1.0 + return {} + + # Create benchmarks dynamically + benchmarks = [] + for prop_name in properties: + benchmarks.append( + MethodBenchmark( + name=prop_name, + method_name=prop_name, + input_generators={"default": gen_mock_data}, + category="property", + ) + ) + + print(f"\nBenchmarking {len(benchmarks)} properties...") + print(f"Config: {config.num_iterations} iterations, {config.warmup_steps} warmup steps") + print(f" {config.num_instances} instances") + print("-" * 80) + + for i, benchmark in enumerate(benchmarks): + + def prop_accessor(prop=benchmark.method_name, **kwargs): + return getattr(rigid_object_data, prop) + + print(f"[{i + 1}/{len(benchmarks)}] [DEFAULT] {benchmark.name}...", end=" ", flush=True) + + result = benchmark_method( + method=prop_accessor, + method_name=benchmark.name, + generator=gen_mock_data, + config=config, + dependencies=PROPERTY_DEPENDENCIES, + ) + result.mode = "default" + results.append(result) + + if result.skipped: + print(f"SKIPPED ({result.skip_reason})") + else: + print(f"{result.mean_time_us:.2f} ± {result.std_time_us:.2f} µs") + + return results + + +def main(): + """Main entry point for the benchmarking script.""" + parser = argparse.ArgumentParser( + description="Micro-benchmarking framework for RigidObjectData class.", + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + ) + parser.add_argument("--num_iterations", type=int, default=1000, help="Number of iterations") + parser.add_argument("--warmup_steps", type=int, default=10, help="Number of warmup steps") + parser.add_argument("--num_instances", type=int, default=4096, help="Number of instances") + parser.add_argument("--device", type=str, default="cuda:0", help="Device") + parser.add_argument("--output", type=str, default=None, help="Output JSON filename") + parser.add_argument("--no_csv", action="store_true", help="Disable CSV output") + + args = parser.parse_args() + + config = BenchmarkConfig( + num_iterations=args.num_iterations, + warmup_steps=args.warmup_steps, + num_instances=args.num_instances, + num_bodies=1, + num_joints=0, + device=args.device, + ) + + results = run_benchmarks(config) + + hardware_info = get_hardware_info() + print_hardware_info(hardware_info) + print_results(results, include_mode=False) + + if args.output: + json_filename = args.output + else: + json_filename = get_default_output_filename("rigid_object_data_benchmark") + + export_results_json(results, config, hardware_info, json_filename) + + if not args.no_csv: + csv_filename = json_filename.replace(".json", ".csv") + export_results_csv(results, csv_filename) + + +if __name__ == "__main__": + main() diff --git a/source/isaaclab_physx/config/extension.toml b/source/isaaclab_physx/config/extension.toml index 2ff0f9bff5a..367b8f23705 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 = "1.0.0" +version = "0.1.3" # 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 4b26e1edc9c..4c292214e6d 100644 --- a/source/isaaclab_physx/docs/CHANGELOG.rst +++ b/source/isaaclab_physx/docs/CHANGELOG.rst @@ -1,7 +1,56 @@ Changelog --------- -1.0.0 (2026-01-28) +0.1.3 (2026-02-03) +~~~~~~~~~~~~~~~~~~~ + +Added +^^^^^ + +* Added :mod:`isaaclab_physx.benchmark` module containing performance micro-benchmarks for + PhysX asset classes. Includes: + + * ``benchmark_articulation.py``: Benchmarks for setter/writer methods on + :class:`~isaaclab_physx.assets.Articulation` including root state, joint state, + joint parameters, and body property operations. + * ``benchmark_articulation_data.py``: Benchmarks for property accessors on + :class:`~isaaclab_physx.assets.ArticulationData` covering root link/COM properties, + joint properties, and body link/COM properties. + * ``benchmark_rigid_object.py``: Benchmarks for setter/writer methods on + :class:`~isaaclab_physx.assets.RigidObject` including root state and body property operations. + * ``benchmark_rigid_object_data.py``: Benchmarks for property accessors on + :class:`~isaaclab_physx.assets.RigidObjectData`. + * ``benchmark_rigid_object_collection.py``: Benchmarks for setter/writer methods on + :class:`~isaaclab_physx.assets.RigidObjectCollection` including body state, pose, + velocity, and property operations. + * ``benchmark_rigid_object_collection_data.py``: Benchmarks for property accessors on + :class:`~isaaclab_physx.assets.RigidObjectCollectionData`. + + All benchmarks support configurable iterations, warmup steps, instance counts, multiple + input modes (torch list, torch tensor), and output to JSON/CSV formats with hardware + information capture. + + +0.1.2 (2026-02-03) +~~~~~~~~~~~~~~~~~~~ + +Added +^^^^^ + +* Added :mod:`isaaclab_physx.test.mock_interfaces` module providing mock PhysX view implementations + for unit testing without requiring Isaac Sim. Includes: + + * :class:`MockRigidBodyView`: Mock for ``physx.RigidBodyView`` with transforms, velocities, + accelerations, and mass properties. + * :class:`MockArticulationView`: Mock for ``physx.ArticulationView`` with root/link states, + DOF properties, and joint control. + * :class:`MockRigidContactView`: Mock for ``physx.RigidContactView`` with contact forces, + positions, normals, and friction data. + * Factory functions including pre-configured quadruped and humanoid views. + * Patching utilities and decorators for easy test injection. + + +0.1.0 (2026-01-28) ~~~~~~~~~~~~~~~~~~ Added diff --git a/source/isaaclab_physx/isaaclab_physx/__init__.py b/source/isaaclab_physx/isaaclab_physx/__init__.py index 69bbf740a59..68a8a553287 100644 --- a/source/isaaclab_physx/isaaclab_physx/__init__.py +++ b/source/isaaclab_physx/isaaclab_physx/__init__.py @@ -17,6 +17,3 @@ # Configure the module-level variables __version__ = ISAACLAB_PHYSX_METADATA["package"]["version"] - -# Import sensors module for auto-registration with factory -from . import sensors # noqa: F401, E402 diff --git a/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py b/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py index ece2816126f..e31f18dccdf 100644 --- a/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py +++ b/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py @@ -16,15 +16,14 @@ import warp as wp from prettytable import PrettyTable -import omni.physics.tensors.impl.api as physx from isaacsim.core.simulation_manager import SimulationManager from pxr import PhysxSchema, UsdPhysics -import isaaclab.sim as sim_utils import isaaclab.utils.math as math_utils -import isaaclab.utils.string as string_utils from isaaclab.actuators import ActuatorBase, ActuatorBaseCfg, ImplicitActuator from isaaclab.assets.articulation.base_articulation import BaseArticulation +from isaaclab.sim.utils.queries import find_first_matching_prim, get_all_matching_child_prims +from isaaclab.utils.string import resolve_matching_names, resolve_matching_names_values from isaaclab.utils.types import ArticulationActions from isaaclab.utils.version import get_isaac_sim_version from isaaclab.utils.wrench_composer import WrenchComposer @@ -32,6 +31,8 @@ from .articulation_data import ArticulationData if TYPE_CHECKING: + import omni.physics.tensors.impl.api as physx + from isaaclab.assets.articulation.articulation_cfg import ArticulationCfg # import logger @@ -166,7 +167,7 @@ def body_names(self) -> list[str]: return self.root_view.shared_metatype.link_names @property - def root_view(self): + def root_view(self) -> physx.ArticulationView: """Root view for the asset. .. note:: @@ -277,7 +278,7 @@ 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 string_utils.resolve_matching_names(name_keys, self.body_names, 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[str] | None = None, preserve_order: bool = False @@ -299,7 +300,7 @@ def find_joints( if joint_subset is None: joint_subset = self.joint_names # find joints - return string_utils.resolve_matching_names(name_keys, joint_subset, preserve_order) + return resolve_matching_names(name_keys, joint_subset, preserve_order) def find_fixed_tendons( self, name_keys: str | Sequence[str], tendon_subsets: list[str] | None = None, preserve_order: bool = False @@ -323,7 +324,7 @@ def find_fixed_tendons( # tendons follow the joint names they are attached to tendon_subsets = self.fixed_tendon_names # find tendons - return string_utils.resolve_matching_names(name_keys, tendon_subsets, preserve_order) + return resolve_matching_names(name_keys, tendon_subsets, preserve_order) def find_spatial_tendons( self, name_keys: str | Sequence[str], tendon_subsets: list[str] | None = None, preserve_order: bool = False @@ -345,7 +346,7 @@ def find_spatial_tendons( if tendon_subsets is None: tendon_subsets = self.spatial_tendon_names # find tendons - return string_utils.resolve_matching_names(name_keys, tendon_subsets, preserve_order) + return resolve_matching_names(name_keys, tendon_subsets, preserve_order) """ Operations - State Writers. @@ -589,6 +590,9 @@ def write_joint_position_to_sim( physx_env_ids = self._ALL_INDICES if joint_ids is None: joint_ids = slice(None) + # convert lists to tensors for proper indexing + if isinstance(env_ids, list): + env_ids = torch.tensor(env_ids, dtype=torch.long, device=self.device) # broadcast env_ids if needed to allow double indexing if env_ids != slice(None) and joint_ids != slice(None): env_ids = env_ids[:, None] @@ -627,6 +631,9 @@ def write_joint_velocity_to_sim( physx_env_ids = self._ALL_INDICES if joint_ids is None: joint_ids = slice(None) + # convert lists to tensors for proper indexing + if isinstance(env_ids, list): + env_ids = torch.tensor(env_ids, dtype=torch.long, device=self.device) # broadcast env_ids if needed to allow double indexing if env_ids != slice(None) and joint_ids != slice(None): env_ids = env_ids[:, None] @@ -662,6 +669,11 @@ def write_joint_stiffness_to_sim( physx_env_ids = self._ALL_INDICES if joint_ids is None: joint_ids = slice(None) + # convert lists to tensors for proper indexing + if isinstance(env_ids, list): + env_ids = torch.tensor(env_ids, dtype=torch.long, device=self.device) + if isinstance(physx_env_ids, list): + physx_env_ids = torch.tensor(physx_env_ids, dtype=torch.long, device=self.device) # broadcast env_ids if needed to allow double indexing if env_ids != slice(None) and joint_ids != slice(None): env_ids = env_ids[:, None] @@ -691,6 +703,11 @@ def write_joint_damping_to_sim( physx_env_ids = self._ALL_INDICES if joint_ids is None: joint_ids = slice(None) + # convert lists to tensors for proper indexing + if isinstance(env_ids, list): + env_ids = torch.tensor(env_ids, dtype=torch.long, device=self.device) + if isinstance(physx_env_ids, list): + physx_env_ids = torch.tensor(physx_env_ids, dtype=torch.long, device=self.device) # broadcast env_ids if needed to allow double indexing if env_ids != slice(None) and joint_ids != slice(None): env_ids = env_ids[:, None] @@ -723,6 +740,11 @@ def write_joint_position_limit_to_sim( physx_env_ids = self._ALL_INDICES if joint_ids is None: joint_ids = slice(None) + # convert lists to tensors for proper indexing + if isinstance(env_ids, list): + env_ids = torch.tensor(env_ids, dtype=torch.long, device=self.device) + if isinstance(physx_env_ids, list): + physx_env_ids = torch.tensor(physx_env_ids, dtype=torch.long, device=self.device) # broadcast env_ids if needed to allow double indexing if env_ids != slice(None) and joint_ids != slice(None): env_ids = env_ids[:, None] @@ -783,6 +805,11 @@ def write_joint_velocity_limit_to_sim( physx_env_ids = self._ALL_INDICES if joint_ids is None: joint_ids = slice(None) + # convert lists to tensors for proper indexing + if isinstance(env_ids, list): + env_ids = torch.tensor(env_ids, dtype=torch.long, device=self.device) + if isinstance(physx_env_ids, list): + physx_env_ids = torch.tensor(physx_env_ids, dtype=torch.long, device=self.device) # broadcast env_ids if needed to allow double indexing if env_ids != slice(None) and joint_ids != slice(None): env_ids = env_ids[:, None] @@ -818,6 +845,11 @@ def write_joint_effort_limit_to_sim( physx_env_ids = self._ALL_INDICES if joint_ids is None: joint_ids = slice(None) + # convert lists to tensors for proper indexing + if isinstance(env_ids, list): + env_ids = torch.tensor(env_ids, dtype=torch.long, device=self.device) + if isinstance(physx_env_ids, list): + physx_env_ids = torch.tensor(physx_env_ids, dtype=torch.long, device=self.device) # broadcast env_ids if needed to allow double indexing if env_ids != slice(None) and joint_ids != slice(None): env_ids = env_ids[:, None] @@ -852,6 +884,11 @@ def write_joint_armature_to_sim( physx_env_ids = self._ALL_INDICES if joint_ids is None: joint_ids = slice(None) + # convert lists to tensors for proper indexing + if isinstance(env_ids, list): + env_ids = torch.tensor(env_ids, dtype=torch.long, device=self.device) + if isinstance(physx_env_ids, list): + physx_env_ids = torch.tensor(physx_env_ids, dtype=torch.long, device=self.device) # broadcast env_ids if needed to allow double indexing if env_ids != slice(None) and joint_ids != slice(None): env_ids = env_ids[:, None] @@ -898,6 +935,11 @@ def write_joint_friction_coefficient_to_sim( physx_env_ids = self._ALL_INDICES if joint_ids is None: joint_ids = slice(None) + # convert lists to tensors for proper indexing + if isinstance(env_ids, list): + env_ids = torch.tensor(env_ids, dtype=torch.long, device=self.device) + if isinstance(physx_env_ids, list): + physx_env_ids = torch.tensor(physx_env_ids, dtype=torch.long, device=self.device) # broadcast env_ids if needed to allow double indexing if env_ids != slice(None) and joint_ids != slice(None): env_ids = env_ids[:, None] @@ -960,6 +1002,11 @@ def write_joint_dynamic_friction_coefficient_to_sim( physx_env_ids = self._ALL_INDICES if joint_ids is None: joint_ids = slice(None) + # convert lists to tensors for proper indexing + if isinstance(env_ids, list): + env_ids = torch.tensor(env_ids, dtype=torch.long, device=self.device) + if isinstance(physx_env_ids, list): + physx_env_ids = torch.tensor(physx_env_ids, dtype=torch.long, device=self.device) # broadcast env_ids if needed to allow double indexing if env_ids != slice(None) and joint_ids != slice(None): env_ids = env_ids[:, None] @@ -994,6 +1041,11 @@ def write_joint_viscous_friction_coefficient_to_sim( physx_env_ids = self._ALL_INDICES if joint_ids is None: joint_ids = slice(None) + # convert lists to tensors for proper indexing + if isinstance(env_ids, list): + env_ids = torch.tensor(env_ids, dtype=torch.long, device=self.device) + if isinstance(physx_env_ids, list): + physx_env_ids = torch.tensor(physx_env_ids, dtype=torch.long, device=self.device) # broadcast env_ids if needed to allow double indexing if env_ids != slice(None) and joint_ids != slice(None): env_ids = env_ids[:, None] @@ -1028,6 +1080,11 @@ def set_masses( physx_env_ids = self._ALL_INDICES if body_ids is None: body_ids = slice(None) + # convert lists to tensors for proper indexing + if isinstance(env_ids, list): + env_ids = torch.tensor(env_ids, dtype=torch.long, device=self.device) + if isinstance(physx_env_ids, list): + physx_env_ids = torch.tensor(physx_env_ids, dtype=torch.long, device=self.device) # broadcast env_ids if needed to allow double indexing if env_ids != slice(None) and body_ids != slice(None): env_ids = env_ids[:, None] @@ -1056,6 +1113,11 @@ def set_coms( physx_env_ids = self._ALL_INDICES if body_ids is None: body_ids = slice(None) + # convert lists to tensors for proper indexing + if isinstance(env_ids, list): + env_ids = torch.tensor(env_ids, dtype=torch.long, device=self.device) + if isinstance(physx_env_ids, list): + physx_env_ids = torch.tensor(physx_env_ids, dtype=torch.long, device=self.device) # broadcast env_ids if needed to allow double indexing if env_ids != slice(None) and body_ids != slice(None): env_ids = env_ids[:, None] @@ -1084,6 +1146,11 @@ def set_inertias( physx_env_ids = self._ALL_INDICES if body_ids is None: body_ids = slice(None) + # convert lists to tensors for proper indexing + if isinstance(env_ids, list): + env_ids = torch.tensor(env_ids, dtype=torch.long, device=self.device) + if isinstance(physx_env_ids, list): + physx_env_ids = torch.tensor(physx_env_ids, dtype=torch.long, device=self.device) # broadcast env_ids if needed to allow double indexing if env_ids != slice(None) and body_ids != slice(None): env_ids = env_ids[:, None] @@ -1201,6 +1268,9 @@ def set_joint_position_target( env_ids = slice(None) if joint_ids is None: joint_ids = slice(None) + # convert lists to tensors for proper indexing + if isinstance(env_ids, list): + env_ids = torch.tensor(env_ids, dtype=torch.long, device=self.device) # broadcast env_ids if needed to allow double indexing if env_ids != slice(None) and joint_ids != slice(None): env_ids = env_ids[:, None] @@ -1228,6 +1298,9 @@ def set_joint_velocity_target( env_ids = slice(None) if joint_ids is None: joint_ids = slice(None) + # convert lists to tensors for proper indexing + if isinstance(env_ids, list): + env_ids = torch.tensor(env_ids, dtype=torch.long, device=self.device) # broadcast env_ids if needed to allow double indexing if env_ids != slice(None) and joint_ids != slice(None): env_ids = env_ids[:, None] @@ -1255,6 +1328,9 @@ def set_joint_effort_target( env_ids = slice(None) if joint_ids is None: joint_ids = slice(None) + # convert lists to tensors for proper indexing + if isinstance(env_ids, list): + env_ids = torch.tensor(env_ids, dtype=torch.long, device=self.device) # broadcast env_ids if needed to allow double indexing if env_ids != slice(None) and joint_ids != slice(None): env_ids = env_ids[:, None] @@ -1287,6 +1363,9 @@ def set_fixed_tendon_stiffness( env_ids = slice(None) if fixed_tendon_ids is None: fixed_tendon_ids = slice(None) + # convert lists to tensors for proper indexing + if isinstance(env_ids, list): + env_ids = torch.tensor(env_ids, dtype=torch.long, device=self.device) if env_ids != slice(None) and fixed_tendon_ids != slice(None): env_ids = env_ids[:, None] # set stiffness @@ -1313,6 +1392,9 @@ def set_fixed_tendon_damping( env_ids = slice(None) if fixed_tendon_ids is None: fixed_tendon_ids = slice(None) + # convert lists to tensors for proper indexing + if isinstance(env_ids, list): + env_ids = torch.tensor(env_ids, dtype=torch.long, device=self.device) if env_ids != slice(None) and fixed_tendon_ids != slice(None): env_ids = env_ids[:, None] # set damping @@ -1340,6 +1422,9 @@ def set_fixed_tendon_limit_stiffness( env_ids = slice(None) if fixed_tendon_ids is None: fixed_tendon_ids = slice(None) + # convert lists to tensors for proper indexing + if isinstance(env_ids, list): + env_ids = torch.tensor(env_ids, dtype=torch.long, device=self.device) if env_ids != slice(None) and fixed_tendon_ids != slice(None): env_ids = env_ids[:, None] # set limit_stiffness @@ -1366,6 +1451,9 @@ def set_fixed_tendon_position_limit( env_ids = slice(None) if fixed_tendon_ids is None: fixed_tendon_ids = slice(None) + # convert lists to tensors for proper indexing + if isinstance(env_ids, list): + env_ids = torch.tensor(env_ids, dtype=torch.long, device=self.device) if env_ids != slice(None) and fixed_tendon_ids != slice(None): env_ids = env_ids[:, None] # set limit @@ -1393,6 +1481,9 @@ def set_fixed_tendon_rest_length( env_ids = slice(None) if fixed_tendon_ids is None: fixed_tendon_ids = slice(None) + # convert lists to tensors for proper indexing + if isinstance(env_ids, list): + env_ids = torch.tensor(env_ids, dtype=torch.long, device=self.device) if env_ids != slice(None) and fixed_tendon_ids != slice(None): env_ids = env_ids[:, None] # set rest_length @@ -1419,6 +1510,9 @@ def set_fixed_tendon_offset( env_ids = slice(None) if fixed_tendon_ids is None: fixed_tendon_ids = slice(None) + # convert lists to tensors for proper indexing + if isinstance(env_ids, list): + env_ids = torch.tensor(env_ids, dtype=torch.long, device=self.device) if env_ids != slice(None) and fixed_tendon_ids != slice(None): env_ids = env_ids[:, None] # set offset @@ -1475,6 +1569,9 @@ def set_spatial_tendon_stiffness( env_ids = slice(None) if spatial_tendon_ids is None: spatial_tendon_ids = slice(None) + # convert lists to tensors for proper indexing + if isinstance(env_ids, list): + env_ids = torch.tensor(env_ids, dtype=torch.long, device=self.device) if env_ids != slice(None) and spatial_tendon_ids != slice(None): env_ids = env_ids[:, None] # set stiffness @@ -1504,6 +1601,9 @@ def set_spatial_tendon_damping( env_ids = slice(None) if spatial_tendon_ids is None: spatial_tendon_ids = slice(None) + # convert lists to tensors for proper indexing + if isinstance(env_ids, list): + env_ids = torch.tensor(env_ids, dtype=torch.long, device=self.device) if env_ids != slice(None) and spatial_tendon_ids != slice(None): env_ids = env_ids[:, None] # set damping @@ -1532,6 +1632,9 @@ def set_spatial_tendon_limit_stiffness( env_ids = slice(None) if spatial_tendon_ids is None: spatial_tendon_ids = slice(None) + # convert lists to tensors for proper indexing + if isinstance(env_ids, list): + env_ids = torch.tensor(env_ids, dtype=torch.long, device=self.device) if env_ids != slice(None) and spatial_tendon_ids != slice(None): env_ids = env_ids[:, None] # set limit stiffness @@ -1559,6 +1662,9 @@ def set_spatial_tendon_offset( env_ids = slice(None) if spatial_tendon_ids is None: spatial_tendon_ids = slice(None) + # convert lists to tensors for proper indexing + if isinstance(env_ids, list): + env_ids = torch.tensor(env_ids, dtype=torch.long, device=self.device) if env_ids != slice(None) and spatial_tendon_ids != slice(None): env_ids = env_ids[:, None] # set offset @@ -1608,13 +1714,13 @@ def _initialize_impl(self): # No articulation root prim path was specified, so we need to search # for it. We search for this in the first environment and then # create a regex that matches all environments. - first_env_matching_prim = sim_utils.find_first_matching_prim(self.cfg.prim_path) + first_env_matching_prim = find_first_matching_prim(self.cfg.prim_path) if first_env_matching_prim is None: raise RuntimeError(f"Failed to find prim for expression: '{self.cfg.prim_path}'.") first_env_matching_prim_path = first_env_matching_prim.GetPath().pathString # Find all articulation root prims in the first environment. - first_env_root_prims = sim_utils.get_all_matching_child_prims( + first_env_root_prims = get_all_matching_child_prims( first_env_matching_prim_path, predicate=lambda prim: prim.HasAPI(UsdPhysics.ArticulationRootAPI), traverse_instance_prims=False, @@ -1715,14 +1821,10 @@ def _process_cfg(self): # -- joint state # joint pos - indices_list, _, values_list = string_utils.resolve_matching_names_values( - self.cfg.init_state.joint_pos, self.joint_names - ) + indices_list, _, values_list = resolve_matching_names_values(self.cfg.init_state.joint_pos, self.joint_names) self.data.default_joint_pos[:, indices_list] = torch.tensor(values_list, device=self.device) # joint vel - indices_list, _, values_list = string_utils.resolve_matching_names_values( - self.cfg.init_state.joint_vel, self.joint_names - ) + indices_list, _, values_list = resolve_matching_names_values(self.cfg.init_state.joint_vel, self.joint_names) self.data.default_joint_vel[:, indices_list] = torch.tensor(values_list, device=self.device) """ diff --git a/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation_data.py b/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation_data.py index 5c648d23287..01300e25f7c 100644 --- a/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation_data.py +++ b/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation_data.py @@ -2,19 +2,23 @@ # All rights reserved. # # SPDX-License-Identifier: BSD-3-Clause +from __future__ import annotations import logging import warnings import weakref +from typing import TYPE_CHECKING import torch -import omni.physics.tensors.impl.api as physx from isaacsim.core.simulation_manager import SimulationManager -import isaaclab.utils.math as math_utils from isaaclab.assets.articulation.base_articulation_data import BaseArticulationData from isaaclab.utils.buffers import TimestampedBuffer +from isaaclab.utils.math import combine_frame_transforms, normalize, quat_apply, quat_apply_inverse + +if TYPE_CHECKING: + from isaaclab.assets.articulation.articulation_view import ArticulationView # import logger logger = logging.getLogger(__name__) @@ -41,7 +45,7 @@ class ArticulationData(BaseArticulationData): __backend_name__: str = "physx" """The name of the backend for the articulation data.""" - def __init__(self, root_view: physx.ArticulationView, device: str): + def __init__(self, root_view: ArticulationView, device: str): """Initializes the articulation data. Args: @@ -52,7 +56,7 @@ def __init__(self, root_view: physx.ArticulationView, device: str): # Set the root articulation view # note: this is stored as a weak reference to avoid circular references between the asset class # and the data container. This is important to avoid memory leaks. - self._root_view: physx.ArticulationView = weakref.proxy(root_view) + self._root_view: ArticulationView = weakref.proxy(root_view) # Set initial time stamp self._sim_timestamp = 0.0 @@ -63,7 +67,7 @@ def __init__(self, root_view: physx.ArticulationView, device: str): gravity = self._physics_sim_view.get_gravity() # Convert to direction vector gravity_dir = torch.tensor((gravity[0], gravity[1], gravity[2]), device=self.device) - gravity_dir = math_utils.normalize(gravity_dir.unsqueeze(0)).squeeze(0) + gravity_dir = normalize(gravity_dir.unsqueeze(0)).squeeze(0) # Initialize constants self.GRAVITY_VEC_W = gravity_dir.repeat(self._root_view.count, 1) @@ -496,7 +500,7 @@ def root_link_vel_w(self) -> torch.Tensor: vel = self.root_com_vel_w.clone() # adjust linear velocity to link from center of mass vel[:, :3] += torch.linalg.cross( - vel[:, 3:], math_utils.quat_apply(self.root_link_quat_w, -self.body_com_pos_b[:, 0]), dim=-1 + vel[:, 3:], quat_apply(self.root_link_quat_w, -self.body_com_pos_b[:, 0]), dim=-1 ) # set the buffer data and timestamp self._root_link_vel_w.data = vel @@ -513,7 +517,7 @@ def root_com_pose_w(self) -> torch.Tensor: """ if self._root_com_pose_w.timestamp < self._sim_timestamp: # apply local transform to center of mass frame - pos, quat = math_utils.combine_frame_transforms( + pos, quat = combine_frame_transforms( self.root_link_pos_w, self.root_link_quat_w, self.body_com_pos_b[:, 0], self.body_com_quat_b[:, 0] ) # set the buffer data and timestamp @@ -621,7 +625,7 @@ def body_link_vel_w(self) -> torch.Tensor: velocities = self.body_com_vel_w.clone() # adjust linear velocity to link from center of mass velocities[..., :3] += torch.linalg.cross( - velocities[..., 3:], math_utils.quat_apply(self.body_link_quat_w, -self.body_com_pos_b), dim=-1 + velocities[..., 3:], quat_apply(self.body_link_quat_w, -self.body_com_pos_b), dim=-1 ) # set the buffer data and timestamp self._body_link_vel_w.data = velocities @@ -639,7 +643,7 @@ def body_com_pose_w(self) -> torch.Tensor: """ if self._body_com_pose_w.timestamp < self._sim_timestamp: # apply local transform to center of mass frame - pos, quat = math_utils.combine_frame_transforms( + pos, quat = combine_frame_transforms( self.body_link_pos_w, self.body_link_quat_w, self.body_com_pos_b, self.body_com_quat_b ) # set the buffer data and timestamp @@ -795,7 +799,7 @@ def joint_acc(self) -> torch.Tensor: @property def projected_gravity_b(self): """Projection of the gravity direction on base frame. Shape is (num_instances, 3).""" - return math_utils.quat_apply_inverse(self.root_link_quat_w, self.GRAVITY_VEC_W) + return quat_apply_inverse(self.root_link_quat_w, self.GRAVITY_VEC_W) @property def heading_w(self): @@ -805,7 +809,7 @@ def heading_w(self): This quantity is computed by assuming that the forward-direction of the base frame is along x-direction, i.e. :math:`(1, 0, 0)`. """ - forward_w = math_utils.quat_apply(self.root_link_quat_w, self.FORWARD_VEC_B) + forward_w = quat_apply(self.root_link_quat_w, self.FORWARD_VEC_B) return torch.atan2(forward_w[:, 1], forward_w[:, 0]) @property @@ -815,7 +819,7 @@ def root_link_lin_vel_b(self) -> torch.Tensor: This quantity is the linear velocity of the articulation root's actor frame with respect to the its actor frame. """ - return math_utils.quat_apply_inverse(self.root_link_quat_w, self.root_link_lin_vel_w) + return quat_apply_inverse(self.root_link_quat_w, self.root_link_lin_vel_w) @property def root_link_ang_vel_b(self) -> torch.Tensor: @@ -824,7 +828,7 @@ def root_link_ang_vel_b(self) -> torch.Tensor: This quantity is the angular velocity of the articulation root's actor frame with respect to the its actor frame. """ - return math_utils.quat_apply_inverse(self.root_link_quat_w, self.root_link_ang_vel_w) + return quat_apply_inverse(self.root_link_quat_w, self.root_link_ang_vel_w) @property def root_com_lin_vel_b(self) -> torch.Tensor: @@ -833,7 +837,7 @@ def root_com_lin_vel_b(self) -> torch.Tensor: This quantity is the linear velocity of the articulation root's center of mass frame with respect to the its actor frame. """ - return math_utils.quat_apply_inverse(self.root_link_quat_w, self.root_com_lin_vel_w) + return quat_apply_inverse(self.root_link_quat_w, self.root_com_lin_vel_w) @property def root_com_ang_vel_b(self) -> torch.Tensor: @@ -842,7 +846,7 @@ def root_com_ang_vel_b(self) -> torch.Tensor: This quantity is the angular velocity of the articulation root's center of mass frame with respect to the its actor frame. """ - return math_utils.quat_apply_inverse(self.root_link_quat_w, self.root_com_ang_vel_w) + return quat_apply_inverse(self.root_link_quat_w, self.root_com_ang_vel_w) """ Sliced properties. 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 fd6c5befeb3..52df1a9d922 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 @@ -428,15 +428,14 @@ def set_masses( if env_ids is None: env_ids = slice(None) physx_env_ids = self._ALL_INDICES - if body_ids is None: - body_ids = slice(None) - # broadcast env_ids if needed to allow double indexing - if env_ids != slice(None) and body_ids != slice(None): - env_ids = env_ids[:, None] + # convert lists to tensors for proper indexing + if isinstance(physx_env_ids, list): + physx_env_ids = torch.tensor(physx_env_ids, dtype=torch.long, device=self.device) # set into internal buffers - self.data.body_mass[env_ids, body_ids] = masses + # _body_mass shape from view is (N, 1), masses input shape is (N, B) where B=1 for RigidObject + self.data._body_mass[env_ids] = masses.reshape(-1, 1) # set into simulation - self.root_view.set_masses(self.data.body_mass.cpu(), indices=physx_env_ids.cpu()) + self.root_view.set_masses(self.data._body_mass.cpu(), indices=physx_env_ids.cpu()) def set_coms( self, @@ -457,15 +456,18 @@ def set_coms( if env_ids is None: env_ids = slice(None) physx_env_ids = self._ALL_INDICES - if body_ids is None: - body_ids = slice(None) - # broadcast env_ids if needed to allow double indexing - if env_ids != slice(None) and body_ids != slice(None): - env_ids = env_ids[:, None] + # Body IDs are always all bodies for a single rigid object. + body_ids = slice(None) + # convert lists to tensors for proper indexing + if isinstance(physx_env_ids, list): + physx_env_ids = torch.tensor(physx_env_ids, dtype=torch.long, device=self.device) # set into internal buffers - self.data.body_com_pose_b[env_ids, body_ids] = coms + # body_com_pose_b shape is (N, B, 7) where [:,:,:3] is position and [:,:,3:7] is quaternion + # coms input shape is (N, B, 3) - only setting the position part + self.data.body_com_pose_b[env_ids, body_ids, :3] = coms # set into simulation - self.root_view.set_coms(self.data.body_com_pose_b.cpu(), indices=physx_env_ids.cpu()) + # RigidBodyView expects (N, 7) but body_com_pose_b is (N, 1, 7), squeeze the body dim + self.root_view.set_coms(self.data.body_com_pose_b.squeeze(1).cpu(), indices=physx_env_ids.cpu()) def set_inertias( self, @@ -485,15 +487,15 @@ def set_inertias( if env_ids is None: env_ids = slice(None) physx_env_ids = self._ALL_INDICES - if body_ids is None: - body_ids = slice(None) - # broadcast env_ids if needed to allow double indexing - if env_ids != slice(None) and body_ids != slice(None): - env_ids = env_ids[:, None] + # convert lists to tensors for proper indexing + if isinstance(physx_env_ids, list): + physx_env_ids = torch.tensor(physx_env_ids, dtype=torch.long, device=self.device) # set into internal buffers - self.data.body_inertia[env_ids, body_ids] = inertias + # _body_inertia shape from view is (N, 9) - flattened 3x3 matrices + # inertias input shape is (N, B, 3, 3) where B=1 for RigidObject, flatten to (N, 9) + self.data._body_inertia[env_ids] = inertias.reshape(-1, 9) # set into simulation - self.root_view.set_inertias(self.data.body_inertia.cpu(), indices=physx_env_ids.cpu()) + self.root_view.set_inertias(self.data._body_inertia.cpu(), indices=physx_env_ids.cpu()) def set_external_force_and_torque( self, diff --git a/source/isaaclab_physx/isaaclab_physx/assets/rigid_object/rigid_object_data.py b/source/isaaclab_physx/isaaclab_physx/assets/rigid_object/rigid_object_data.py index ce7aa0b4b2a..cc3e636aa04 100644 --- a/source/isaaclab_physx/isaaclab_physx/assets/rigid_object/rigid_object_data.py +++ b/source/isaaclab_physx/isaaclab_physx/assets/rigid_object/rigid_object_data.py @@ -2,19 +2,28 @@ # All rights reserved. # # SPDX-License-Identifier: BSD-3-Clause +from __future__ import annotations import logging import warnings import weakref +from typing import TYPE_CHECKING import torch -import omni.physics.tensors.impl.api as physx +from isaacsim.core.simulation_manager import SimulationManager -import isaaclab.utils.math as math_utils from isaaclab.assets.rigid_object.base_rigid_object_data import BaseRigidObjectData -from isaaclab.sim.utils.stage import get_current_stage_id from isaaclab.utils.buffers import TimestampedBuffer +from isaaclab.utils.math import ( + combine_frame_transforms, + normalize, + quat_apply, + quat_apply_inverse, +) + +if TYPE_CHECKING: + from isaaclab.assets.rigid_object.rigid_object_view import RigidObjectView # import logger logger = logging.getLogger(__name__) @@ -44,7 +53,7 @@ class RigidObjectData(BaseRigidObjectData): __backend_name__: str = "physx" """The name of the backend for the rigid object data.""" - def __init__(self, root_view: physx.RigidBodyView, device: str): + def __init__(self, root_view: RigidObjectView, device: str): """Initializes the rigid object data. Args: @@ -55,20 +64,18 @@ def __init__(self, root_view: physx.RigidBodyView, device: str): # Set the root rigid body view # note: this is stored as a weak reference to avoid circular references between the asset class # and the data container. This is important to avoid memory leaks. - self._root_view: physx.RigidBodyView = weakref.proxy(root_view) + self._root_view: RigidObjectView = weakref.proxy(root_view) # Set initial time stamp self._sim_timestamp = 0.0 self._is_primed = False # Obtain global physics sim view - stage_id = get_current_stage_id() - physics_sim_view = physx.create_simulation_view("torch", stage_id) - physics_sim_view.set_subspace_roots("/") - gravity = physics_sim_view.get_gravity() + self._physics_sim_view = SimulationManager.get_physics_sim_view() + gravity = self._physics_sim_view.get_gravity() # Convert to direction vector gravity_dir = torch.tensor((gravity[0], gravity[1], gravity[2]), device=self.device) - gravity_dir = math_utils.normalize(gravity_dir.unsqueeze(0)).squeeze(0) + gravity_dir = normalize(gravity_dir.unsqueeze(0)).squeeze(0) # Initialize constants self.GRAVITY_VEC_W = gravity_dir.repeat(self._root_view.count, 1) @@ -229,7 +236,7 @@ def root_link_vel_w(self) -> torch.Tensor: vel = self.root_com_vel_w.clone() # adjust linear velocity to link from center of mass vel[:, :3] += torch.linalg.cross( - vel[:, 3:], math_utils.quat_apply(self.root_link_quat_w, -self.body_com_pos_b[:, 0]), dim=-1 + vel[:, 3:], quat_apply(self.root_link_quat_w, -self.body_com_pos_b[:, 0]), dim=-1 ) # set the buffer data and timestamp self._root_link_vel_w.data = vel @@ -246,7 +253,7 @@ def root_com_pose_w(self) -> torch.Tensor: """ if self._root_com_pose_w.timestamp < self._sim_timestamp: # apply local transform to center of mass frame - pos, quat = math_utils.combine_frame_transforms( + pos, quat = combine_frame_transforms( self.root_link_pos_w, self.root_link_quat_w, self.body_com_pos_b[:, 0], self.body_com_quat_b[:, 0] ) # set the buffer data and timestamp @@ -427,7 +434,7 @@ def body_com_pose_b(self) -> torch.Tensor: @property def projected_gravity_b(self) -> torch.Tensor: """Projection of the gravity direction on base frame. Shape is (num_instances, 3).""" - return math_utils.quat_apply_inverse(self.root_link_quat_w, self.GRAVITY_VEC_W) + return quat_apply_inverse(self.root_link_quat_w, self.GRAVITY_VEC_W) @property def heading_w(self) -> torch.Tensor: @@ -437,7 +444,7 @@ def heading_w(self) -> torch.Tensor: This quantity is computed by assuming that the forward-direction of the base frame is along x-direction, i.e. :math:`(1, 0, 0)`. """ - forward_w = math_utils.quat_apply(self.root_link_quat_w, self.FORWARD_VEC_B) + forward_w = quat_apply(self.root_link_quat_w, self.FORWARD_VEC_B) return torch.atan2(forward_w[:, 1], forward_w[:, 0]) @property @@ -447,7 +454,7 @@ def root_link_lin_vel_b(self) -> torch.Tensor: This quantity is the linear velocity of the actor frame of the root rigid body frame with respect to the rigid body's actor frame. """ - return math_utils.quat_apply_inverse(self.root_link_quat_w, self.root_link_lin_vel_w) + return quat_apply_inverse(self.root_link_quat_w, self.root_link_lin_vel_w) @property def root_link_ang_vel_b(self) -> torch.Tensor: @@ -456,7 +463,7 @@ def root_link_ang_vel_b(self) -> torch.Tensor: This quantity is the angular velocity of the actor frame of the root rigid body frame with respect to the rigid body's actor frame. """ - return math_utils.quat_apply_inverse(self.root_link_quat_w, self.root_link_ang_vel_w) + return quat_apply_inverse(self.root_link_quat_w, self.root_link_ang_vel_w) @property def root_com_lin_vel_b(self) -> torch.Tensor: @@ -465,7 +472,7 @@ def root_com_lin_vel_b(self) -> torch.Tensor: This quantity is the linear velocity of the root rigid body's center of mass frame with respect to the rigid body's actor frame. """ - return math_utils.quat_apply_inverse(self.root_link_quat_w, self.root_com_lin_vel_w) + return quat_apply_inverse(self.root_link_quat_w, self.root_com_lin_vel_w) @property def root_com_ang_vel_b(self) -> torch.Tensor: @@ -474,7 +481,7 @@ def root_com_ang_vel_b(self) -> torch.Tensor: This quantity is the angular velocity of the root rigid body's center of mass frame with respect to the rigid body's actor frame. """ - return math_utils.quat_apply_inverse(self.root_link_quat_w, self.root_com_ang_vel_w) + return quat_apply_inverse(self.root_link_quat_w, self.root_com_ang_vel_w) """ Sliced properties. @@ -666,10 +673,6 @@ def _create_buffers(self) -> None: self._body_mass = self._root_view.get_masses().to(self.device).clone() self._body_inertia = self._root_view.get_inertias().to(self.device).clone() - # -- Default mass and inertia (Lazy allocation of default values) - self._default_mass = None - self._default_inertia = None - """ Backwards compatibility. (Deprecated properties) """ 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 acc0db866f1..7223aa4b94c 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 @@ -331,6 +331,12 @@ def write_body_link_pose_to_sim( if body_ids is None: body_ids = self._ALL_BODY_INDICES + # convert lists to tensors for proper indexing + if isinstance(env_ids, list): + env_ids = torch.tensor(env_ids, dtype=torch.long, device=self.device) + if isinstance(body_ids, list): + body_ids = torch.tensor(body_ids, dtype=torch.long, device=self.device) + # note: we need to do this here since tensors are not set into simulation until step. # set into internal buffers self.data.body_link_pose_w[env_ids[:, None], body_ids] = body_poses.clone() @@ -382,6 +388,12 @@ def write_body_com_pose_to_sim( if body_ids is None: body_ids = self._ALL_BODY_INDICES + # convert lists to tensors for proper indexing + if isinstance(env_ids, list): + env_ids = torch.tensor(env_ids, dtype=torch.long, device=self.device) + if isinstance(body_ids, list): + body_ids = torch.tensor(body_ids, dtype=torch.long, device=self.device) + # set into internal buffers self.data.body_com_pose_w[env_ids[:, None], body_ids] = body_poses.clone() # update these buffers only if the user is using them. Otherwise this adds to overhead. @@ -446,6 +458,12 @@ def write_body_com_velocity_to_sim( if body_ids is None: body_ids = self._ALL_BODY_INDICES + # convert lists to tensors for proper indexing + if isinstance(env_ids, list): + env_ids = torch.tensor(env_ids, dtype=torch.long, device=self.device) + if isinstance(body_ids, list): + body_ids = torch.tensor(body_ids, dtype=torch.long, device=self.device) + # note: we need to do this here since tensors are not set into simulation until step. # set into internal buffers self.data.body_com_vel_w[env_ids[:, None], body_ids] = body_velocities.clone() @@ -487,6 +505,12 @@ def write_body_link_velocity_to_sim( if body_ids is None: body_ids = self._ALL_BODY_INDICES + # convert lists to tensors for proper indexing + if isinstance(env_ids, list): + env_ids = torch.tensor(env_ids, dtype=torch.long, device=self.device) + if isinstance(body_ids, list): + body_ids = torch.tensor(body_ids, dtype=torch.long, device=self.device) + # set into internal buffers self.data.body_link_vel_w[env_ids[:, None], body_ids] = body_velocities.clone() # update these buffers only if the user is using them. Otherwise this adds to overhead. @@ -518,10 +542,31 @@ def set_masses( """Set masses of all bodies. Args: - masses: Masses of all bodies. Shape is (num_instances, num_bodies). + masses: Masses of all bodies. Shape is (len(env_ids), len(body_ids)). body_ids: The body indices to set the masses for. Defaults to None (all bodies). env_ids: The environment indices to set the masses for. Defaults to None (all environments). """ + # resolve indices + if env_ids is None: + env_ids = self._ALL_ENV_INDICES + if body_ids is None: + body_ids = self._ALL_BODY_INDICES + + # convert lists to tensors for proper indexing + if isinstance(env_ids, list): + env_ids = torch.tensor(env_ids, dtype=torch.long, device=self.device) + if isinstance(body_ids, list): + body_ids = torch.tensor(body_ids, dtype=torch.long, device=self.device) + + # set into internal buffers + # _body_mass shape from view is (num_instances * num_bodies, 1) + # We need to update only the selected env_ids and body_ids + view_ids = self._env_body_ids_to_view_ids(env_ids, body_ids) + # masses input shape is (len(env_ids), len(body_ids)), flatten to match view + self.data._body_mass[view_ids] = masses.reshape(-1, 1) + + # set into simulation + self.root_view.set_masses(self.data._body_mass.cpu(), indices=view_ids.cpu()) def set_coms( self, @@ -532,11 +577,37 @@ def set_coms( """Set center of mass positions of all bodies. Args: - coms: Center of mass positions of all bodies. Shape is (num_instances, num_bodies, 3). + coms: Center of mass positions of all bodies. Shape is (len(env_ids), len(body_ids), 3). body_ids: The body indices to set the center of mass positions for. Defaults to None (all bodies). env_ids: The environment indices to set the center of mass positions for. Defaults to None (all environments). """ + # resolve indices + if env_ids is None: + env_ids = self._ALL_ENV_INDICES + if body_ids is None: + body_ids = self._ALL_BODY_INDICES + + # convert lists to tensors for proper indexing + if isinstance(env_ids, list): + env_ids = torch.tensor(env_ids, dtype=torch.long, device=self.device) + if isinstance(body_ids, list): + body_ids = torch.tensor(body_ids, dtype=torch.long, device=self.device) + + # get view indices + view_ids = self._env_body_ids_to_view_ids(env_ids, body_ids) + + # get current com poses and update position part + # body_com_pose_b triggers lazy evaluation, so we work with the underlying buffer + current_poses = self.root_view.get_coms().to(self.device) + # coms input shape is (len(env_ids), len(body_ids), 3), flatten to (N, 3) + current_poses[view_ids, :3] = coms.reshape(-1, 3) + + # set into simulation + self.root_view.set_coms(current_poses.cpu(), indices=view_ids.cpu()) + + # invalidate the cached buffer + self.data._body_com_pose_b.timestamp = -1 def set_inertias( self, @@ -547,10 +618,32 @@ def set_inertias( """Set inertias of all bodies. Args: - inertias: Inertias of all bodies. Shape is (num_instances, num_bodies, 3, 3). + inertias: Inertias of all bodies. Shape is (len(env_ids), len(body_ids), 3, 3). body_ids: The body indices to set the inertias for. Defaults to None (all bodies). env_ids: The environment indices to set the inertias for. Defaults to None (all environments). """ + # resolve indices + if env_ids is None: + env_ids = self._ALL_ENV_INDICES + if body_ids is None: + body_ids = self._ALL_BODY_INDICES + + # convert lists to tensors for proper indexing + if isinstance(env_ids, list): + env_ids = torch.tensor(env_ids, dtype=torch.long, device=self.device) + if isinstance(body_ids, list): + body_ids = torch.tensor(body_ids, dtype=torch.long, device=self.device) + + # get view indices + view_ids = self._env_body_ids_to_view_ids(env_ids, body_ids) + + # set into internal buffers + # _body_inertia shape from view is (num_instances * num_bodies, 9) - flattened 3x3 matrices + # inertias input shape is (len(env_ids), len(body_ids), 3, 3), flatten to (N, 9) + self.data._body_inertia[view_ids] = inertias.reshape(-1, 9) + + # set into simulation + self.root_view.set_inertias(self.data._body_inertia.cpu(), indices=view_ids.cpu()) def set_external_force_and_torque( self, diff --git a/source/isaaclab_physx/isaaclab_physx/assets/rigid_object_collection/rigid_object_collection_data.py b/source/isaaclab_physx/isaaclab_physx/assets/rigid_object_collection/rigid_object_collection_data.py index 58492572a1c..64acd87ceb1 100644 --- a/source/isaaclab_physx/isaaclab_physx/assets/rigid_object_collection/rigid_object_collection_data.py +++ b/source/isaaclab_physx/isaaclab_physx/assets/rigid_object_collection/rigid_object_collection_data.py @@ -2,19 +2,23 @@ # All rights reserved. # # SPDX-License-Identifier: BSD-3-Clause +from __future__ import annotations import logging import warnings import weakref +from typing import TYPE_CHECKING import torch -import omni.physics.tensors.impl.api as physx +from isaacsim.core.simulation_manager import SimulationManager -import isaaclab.utils.math as math_utils from isaaclab.assets.rigid_object_collection.base_rigid_object_collection_data import BaseRigidObjectCollectionData -from isaaclab.sim.utils.stage import get_current_stage_id from isaaclab.utils.buffers import TimestampedBuffer +from isaaclab.utils.math import combine_frame_transforms, normalize, quat_apply, quat_apply_inverse + +if TYPE_CHECKING: + from isaaclab.assets.rigid_object_collection.rigid_object_collection_view import RigidObjectCollectionView # import logger logger = logging.getLogger(__name__) @@ -44,7 +48,7 @@ class RigidObjectCollectionData(BaseRigidObjectCollectionData): __backend_name__: str = "physx" """The name of the backend for the rigid object collection data.""" - def __init__(self, root_view: physx.RigidBodyView, num_bodies: int, device: str): + def __init__(self, root_view: RigidObjectCollectionView, num_bodies: int, device: str): """Initializes the rigid object data. Args: @@ -57,7 +61,7 @@ def __init__(self, root_view: physx.RigidBodyView, num_bodies: int, device: str) # Set the root rigid body view # note: this is stored as a weak reference to avoid circular references between the asset class # and the data container. This is important to avoid memory leaks. - self._root_view: physx.RigidBodyView = weakref.proxy(root_view) + self._root_view: RigidObjectCollectionView = weakref.proxy(root_view) self.num_instances = self._root_view.count // self.num_bodies # Set initial time stamp @@ -65,13 +69,11 @@ def __init__(self, root_view: physx.RigidBodyView, num_bodies: int, device: str) self._is_primed = False # Obtain global physics sim view - stage_id = get_current_stage_id() - physics_sim_view = physx.create_simulation_view("torch", stage_id) - physics_sim_view.set_subspace_roots("/") - gravity = physics_sim_view.get_gravity() + self._physics_sim_view = SimulationManager.get_physics_sim_view() + gravity = self._physics_sim_view.get_gravity() # Convert to direction vector gravity_dir = torch.tensor((gravity[0], gravity[1], gravity[2]), device=self.device) - gravity_dir = math_utils.normalize(gravity_dir.unsqueeze(0)).squeeze(0) + gravity_dir = normalize(gravity_dir.unsqueeze(0)).squeeze(0) # Initialize constants self.GRAVITY_VEC_W = gravity_dir.repeat(self.num_instances, self.num_bodies, 1) @@ -209,7 +211,7 @@ def default_body_state(self, value: torch.Tensor) -> None: @property def body_link_pose_w(self) -> torch.Tensor: - """Body link pose ``[pos, quat]`` in simulation world frame. Shape is (num_instances, 1, 7). + """Body link pose ``[pos, quat]`` in simulation world frame. Shape is (num_instances, num_bodies, 7). This quantity is the pose of the actor frame of the rigid body relative to the world. The orientation is provided in (x, y, z, w) format. @@ -225,7 +227,7 @@ def body_link_pose_w(self) -> torch.Tensor: @property def body_link_vel_w(self) -> torch.Tensor: - """Body link velocity ``[lin_vel, ang_vel]`` in simulation world frame. Shape is (num_instances, 1, 6). + """Body link velocity ``[lin_vel, ang_vel]`` in simulation world frame. Shape is (num_instances, num_bodies, 6). This quantity contains the linear and angular velocities of the actor frame of the root rigid body relative to the world. @@ -235,7 +237,7 @@ def body_link_vel_w(self) -> torch.Tensor: velocity = self.body_com_vel_w.clone() # adjust linear velocity to link from center of mass velocity[..., :3] += torch.linalg.cross( - velocity[..., 3:], math_utils.quat_apply(self.body_link_quat_w, -self.body_com_pos_b), dim=-1 + velocity[..., 3:], quat_apply(self.body_link_quat_w, -self.body_com_pos_b), dim=-1 ) # set the buffer data and timestamp self._body_link_vel_w.data = velocity @@ -245,14 +247,14 @@ def body_link_vel_w(self) -> torch.Tensor: @property def body_com_pose_w(self) -> torch.Tensor: - """Body center of mass pose ``[pos, quat]`` in simulation world frame. Shape is (num_instances, 1, 7). + """Body center of mass pose ``[pos, quat]`` in simulation world frame. Shape is (num_instances, num_bodies, 7). This quantity is the pose of the center of mass frame of the rigid body relative to the world. The orientation is provided in (x, y, z, w) format. """ if self._body_com_pose_w.timestamp < self._sim_timestamp: # adjust pose to center of mass - pos, quat = math_utils.combine_frame_transforms( + pos, quat = combine_frame_transforms( self.body_link_pos_w, self.body_link_quat_w, self.body_com_pos_b, self.body_com_quat_b ) # set the buffer data and timestamp @@ -264,7 +266,7 @@ def body_com_pose_w(self) -> torch.Tensor: @property def body_com_vel_w(self) -> torch.Tensor: """Body center of mass velocity ``[lin_vel, ang_vel]`` in simulation world frame. - Shape is (num_instances, 1, 6). + Shape is (num_instances, num_bodies, 6). This quantity contains the linear and angular velocities of the root rigid body's center of mass frame relative to the world. @@ -278,7 +280,7 @@ def body_com_vel_w(self) -> torch.Tensor: @property def body_state_w(self) -> torch.Tensor: """State of all bodies `[pos, quat, lin_vel, ang_vel]` in simulation world frame. - Shape is (num_instances, 1, 13). + Shape is (num_instances, num_bodies, 13). The position and orientation are of the rigid bodies' actor frame. Meanwhile, the linear and angular velocities are of the rigid bodies' center of mass frame. @@ -292,7 +294,7 @@ def body_state_w(self) -> torch.Tensor: @property def body_link_state_w(self) -> torch.Tensor: """State of all bodies ``[pos, quat, lin_vel, ang_vel]`` in simulation world frame. - Shape is (num_instances, 1, 13). + Shape is (num_instances, num_bodies, 13). The position, quaternion, and linear/angular velocity are of the body's link frame relative to the world. The orientation is provided in (x, y, z, w) format. @@ -321,7 +323,7 @@ def body_com_state_w(self) -> torch.Tensor: @property def body_com_acc_w(self) -> torch.Tensor: """Acceleration of all bodies ``[lin_acc, ang_acc]`` in the simulation world frame. - Shape is (num_instances, 1, 6). + Shape is (num_instances, num_bodies, 6). This quantity is the acceleration of the rigid bodies' center of mass frame relative to the world. """ @@ -333,7 +335,7 @@ def body_com_acc_w(self) -> torch.Tensor: @property def body_com_pose_b(self) -> torch.Tensor: """Center of mass pose ``[pos, quat]`` of all bodies in their respective body's link frames. - Shape is (num_instances, 1, 7). + Shape is (num_instances, num_bodies, 7). This quantity is the pose of the center of mass frame of the rigid body relative to the body's link frame. The orientation is provided in (x, y, z, w) format. @@ -349,12 +351,12 @@ def body_com_pose_b(self) -> torch.Tensor: @property def body_mass(self) -> torch.Tensor: - """Mass of all bodies in the simulation world frame. Shape is (num_instances, 1, 1).""" + """Mass of all bodies in the simulation world frame. Shape is (num_instances, num_bodies).""" return self._reshape_view_to_data(self._body_mass.to(self.device)) @property def body_inertia(self) -> torch.Tensor: - """Inertia of all bodies in the simulation world frame. Shape is (num_instances, 1, 3, 3).""" + """Inertia of all bodies in the simulation world frame. Shape is (num_instances, num_bodies, 9).""" return self._reshape_view_to_data(self._body_inertia.to(self.device)) """ @@ -363,55 +365,55 @@ def body_inertia(self) -> torch.Tensor: @property def projected_gravity_b(self) -> torch.Tensor: - """Projection of the gravity direction on base frame. Shape is (num_instances, 3).""" - return math_utils.quat_apply_inverse(self.body_link_quat_w, self.GRAVITY_VEC_W) + """Projection of the gravity direction on base frame. Shape is (num_instances, num_bodies, 3).""" + return quat_apply_inverse(self.body_link_quat_w, self.GRAVITY_VEC_W) @property def heading_w(self) -> torch.Tensor: - """Yaw heading of the base frame (in radians). Shape is (num_instances,). + """Yaw heading of the base frame (in radians). Shape is (num_instances, num_bodies). .. note:: This quantity is computed by assuming that the forward-direction of the base frame is along x-direction, i.e. :math:`(1, 0, 0)`. """ - forward_w = math_utils.quat_apply(self.body_link_quat_w, self.FORWARD_VEC_B) + forward_w = quat_apply(self.body_link_quat_w, self.FORWARD_VEC_B) return torch.atan2(forward_w[..., 1], forward_w[..., 0]) @property def body_link_lin_vel_b(self) -> torch.Tensor: - """Root link linear velocity in base frame. Shape is (num_instances, 3). + """Root link linear velocity in base frame. Shape is (num_instances, num_bodies, 3). This quantity is the linear velocity of the actor frame of the root rigid body frame with respect to the rigid body's actor frame. """ - return math_utils.quat_apply_inverse(self.body_link_quat_w, self.body_link_lin_vel_w) + return quat_apply_inverse(self.body_link_quat_w, self.body_link_lin_vel_w) @property def body_link_ang_vel_b(self) -> torch.Tensor: - """Root link angular velocity in base world frame. Shape is (num_instances, 3). + """Root link angular velocity in base world frame. Shape is (num_instances, num_bodies, 3). This quantity is the angular velocity of the actor frame of the root rigid body frame with respect to the rigid body's actor frame. """ - return math_utils.quat_apply_inverse(self.body_link_quat_w, self.body_link_ang_vel_w) + return quat_apply_inverse(self.body_link_quat_w, self.body_link_ang_vel_w) @property def body_com_lin_vel_b(self) -> torch.Tensor: - """Root center of mass linear velocity in base frame. Shape is (num_instances, 3). + """Root center of mass linear velocity in base frame. Shape is (num_instances, num_bodies, 3). This quantity is the linear velocity of the root rigid body's center of mass frame with respect to the rigid body's actor frame. """ - return math_utils.quat_apply_inverse(self.body_link_quat_w, self.body_com_lin_vel_w) + return quat_apply_inverse(self.body_link_quat_w, self.body_com_lin_vel_w) @property def body_com_ang_vel_b(self) -> torch.Tensor: - """Root center of mass angular velocity in base world frame. Shape is (num_instances, 3). + """Root center of mass angular velocity in base world frame. Shape is (num_instances, num_bodies, 3). This quantity is the angular velocity of the root rigid body's center of mass frame with respect to the rigid body's actor frame. """ - return math_utils.quat_apply_inverse(self.body_link_quat_w, self.body_com_ang_vel_w) + return quat_apply_inverse(self.body_link_quat_w, self.body_com_ang_vel_w) """ Sliced properties. @@ -419,7 +421,7 @@ def body_com_ang_vel_b(self) -> torch.Tensor: @property def body_link_pos_w(self) -> torch.Tensor: - """Positions of all bodies in simulation world frame. Shape is (num_instances, 1, 3). + """Positions of all bodies in simulation world frame. Shape is (num_instances, num_bodies, 3). This quantity is the position of the rigid bodies' actor frame relative to the world. """ @@ -427,7 +429,7 @@ def body_link_pos_w(self) -> torch.Tensor: @property def body_link_quat_w(self) -> torch.Tensor: - """Orientation (x, y, z, w) of all bodies in simulation world frame. Shape is (num_instances, 1, 4). + """Orientation (x, y, z, w) of all bodies in simulation world frame. Shape is (num_instances, num_bodies, 4). This quantity is the orientation of the rigid bodies' actor frame relative to the world. """ @@ -435,7 +437,7 @@ def body_link_quat_w(self) -> torch.Tensor: @property def body_link_lin_vel_w(self) -> torch.Tensor: - """Linear velocity of all bodies in simulation world frame. Shape is (num_instances, 1, 3). + """Linear velocity of all bodies in simulation world frame. Shape is (num_instances, num_bodies, 3). This quantity is the linear velocity of the rigid bodies' center of mass frame relative to the world. """ @@ -443,7 +445,7 @@ def body_link_lin_vel_w(self) -> torch.Tensor: @property def body_link_ang_vel_w(self) -> torch.Tensor: - """Angular velocity of all bodies in simulation world frame. Shape is (num_instances, 1, 3). + """Angular velocity of all bodies in simulation world frame. Shape is (num_instances, num_bodies, 3). This quantity is the angular velocity of the rigid bodies' center of mass frame relative to the world. """ @@ -451,7 +453,7 @@ def body_link_ang_vel_w(self) -> torch.Tensor: @property def body_com_pos_w(self) -> torch.Tensor: - """Positions of all bodies in simulation world frame. Shape is (num_instances, 1, 3). + """Positions of all bodies in simulation world frame. Shape is (num_instances, num_bodies, 3). This quantity is the position of the rigid bodies' actor frame. """ @@ -461,13 +463,13 @@ def body_com_pos_w(self) -> torch.Tensor: def body_com_quat_w(self) -> torch.Tensor: """Orientation (x, y, z, w) of the principle axis of inertia of all bodies in simulation world frame. - Shape is (num_instances, 1, 4). This quantity is the orientation of the rigid bodies' actor frame. + Shape is (num_instances, num_bodies, 4). This quantity is the orientation of the rigid bodies' actor frame. """ return self.body_com_pose_w[..., 3:7] @property def body_com_lin_vel_w(self) -> torch.Tensor: - """Linear velocity of all bodies in simulation world frame. Shape is (num_instances, 1, 3). + """Linear velocity of all bodies in simulation world frame. Shape is (num_instances, num_bodies, 3). This quantity is the linear velocity of the rigid bodies' center of mass frame. """ @@ -475,7 +477,7 @@ def body_com_lin_vel_w(self) -> torch.Tensor: @property def body_com_ang_vel_w(self) -> torch.Tensor: - """Angular velocity of all bodies in simulation world frame. Shape is (num_instances, 1, 3). + """Angular velocity of all bodies in simulation world frame. Shape is (num_instances, num_bodies, 3). This quantity is the angular velocity of the rigid bodies' center of mass frame. """ @@ -483,7 +485,7 @@ def body_com_ang_vel_w(self) -> torch.Tensor: @property def body_com_lin_acc_w(self) -> torch.Tensor: - """Linear acceleration of all bodies in simulation world frame. Shape is (num_instances, 1, 3). + """Linear acceleration of all bodies in simulation world frame. Shape is (num_instances, num_bodies, 3). This quantity is the linear acceleration of the rigid bodies' center of mass frame. """ @@ -491,7 +493,7 @@ def body_com_lin_acc_w(self) -> torch.Tensor: @property def body_com_ang_acc_w(self) -> torch.Tensor: - """Angular acceleration of all bodies in simulation world frame. Shape is (num_instances, 1, 3). + """Angular acceleration of all bodies in simulation world frame. Shape is (num_instances, num_bodies, 3). This quantity is the angular acceleration of the rigid bodies' center of mass frame. """ @@ -500,7 +502,7 @@ def body_com_ang_acc_w(self) -> torch.Tensor: @property def body_com_pos_b(self) -> torch.Tensor: """Center of mass position of all of the bodies in their respective link frames. - Shape is (num_instances, 1, 3). + Shape is (num_instances, num_bodies, 3). This quantity is the center of mass location relative to its body'slink frame. """ @@ -509,7 +511,7 @@ def body_com_pos_b(self) -> torch.Tensor: @property def body_com_quat_b(self) -> torch.Tensor: """Orientation (x, y, z, w) of the principle axis of inertia of all of the bodies in their - respective link frames. Shape is (num_instances, 1, 4). + respective link frames. Shape is (num_instances, num_bodies, 4). This quantity is the orientation of the principles axes of inertia relative to its body's link frame. """ diff --git a/source/isaaclab_physx/isaaclab_physx/test/__init__.py b/source/isaaclab_physx/isaaclab_physx/test/__init__.py new file mode 100644 index 00000000000..7d7ed09bafb --- /dev/null +++ b/source/isaaclab_physx/isaaclab_physx/test/__init__.py @@ -0,0 +1,6 @@ +# 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 utilities for isaaclab_physx.""" diff --git a/source/isaaclab_physx/isaaclab_physx/test/mock_interfaces/__init__.py b/source/isaaclab_physx/isaaclab_physx/test/mock_interfaces/__init__.py new file mode 100644 index 00000000000..8e619ad9227 --- /dev/null +++ b/source/isaaclab_physx/isaaclab_physx/test/mock_interfaces/__init__.py @@ -0,0 +1,32 @@ +# 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 + +"""Mock interfaces for PhysX TensorAPI views. + +This module provides mock implementations of PhysX TensorAPI views for unit testing +without requiring Isaac Sim or GPU simulation. +""" + +from .factories import ( + create_mock_articulation_view, + create_mock_humanoid_view, + create_mock_quadruped_view, + create_mock_rigid_body_view, + create_mock_rigid_contact_view, +) +from .views import MockArticulationView, MockRigidBodyView, MockRigidContactView + +__all__ = [ + # Views + "MockRigidBodyView", + "MockArticulationView", + "MockRigidContactView", + # Factories + "create_mock_rigid_body_view", + "create_mock_articulation_view", + "create_mock_rigid_contact_view", + "create_mock_quadruped_view", + "create_mock_humanoid_view", +] diff --git a/source/isaaclab_physx/isaaclab_physx/test/mock_interfaces/factories.py b/source/isaaclab_physx/isaaclab_physx/test/mock_interfaces/factories.py new file mode 100644 index 00000000000..a3d5ae812c9 --- /dev/null +++ b/source/isaaclab_physx/isaaclab_physx/test/mock_interfaces/factories.py @@ -0,0 +1,240 @@ +# 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 + +"""Factory functions for creating mock PhysX views.""" + +from __future__ import annotations + +from .views import MockArticulationView, MockRigidBodyView, MockRigidContactView + + +def create_mock_rigid_body_view( + count: int = 1, + prim_paths: list[str] | None = None, + device: str = "cpu", +) -> MockRigidBodyView: + """Create a mock rigid body view. + + Args: + count: Number of rigid body instances. + prim_paths: USD prim paths for each instance. + device: Device for tensor allocation. + + Returns: + A MockRigidBodyView instance. + """ + return MockRigidBodyView(count=count, prim_paths=prim_paths, device=device) + + +def create_mock_articulation_view( + count: int = 1, + num_dofs: int = 1, + num_links: int = 2, + dof_names: list[str] | None = None, + link_names: list[str] | None = None, + fixed_base: bool = False, + prim_paths: list[str] | None = None, + device: str = "cpu", +) -> MockArticulationView: + """Create a mock articulation view. + + Args: + count: Number of articulation instances. + num_dofs: Number of degrees of freedom (joints). + num_links: Number of links (bodies). + dof_names: Names of the DOFs. + link_names: Names of the links. + fixed_base: Whether the articulation has a fixed base. + prim_paths: USD prim paths for each instance. + device: Device for tensor allocation. + + Returns: + A MockArticulationView instance. + """ + return MockArticulationView( + count=count, + num_dofs=num_dofs, + num_links=num_links, + dof_names=dof_names, + link_names=link_names, + fixed_base=fixed_base, + prim_paths=prim_paths, + device=device, + ) + + +def create_mock_rigid_contact_view( + count: int = 1, + num_bodies: int = 1, + filter_count: int = 0, + max_contact_data_count: int = 16, + device: str = "cpu", +) -> MockRigidContactView: + """Create a mock rigid contact view. + + Args: + count: Number of instances. + num_bodies: Number of bodies per instance. + filter_count: Number of filter bodies for contact filtering. + max_contact_data_count: Maximum number of contact data points. + device: Device for tensor allocation. + + Returns: + A MockRigidContactView instance. + """ + return MockRigidContactView( + count=count, + num_bodies=num_bodies, + filter_count=filter_count, + max_contact_data_count=max_contact_data_count, + device=device, + ) + + +# -- Pre-configured factories -- + + +def create_mock_quadruped_view( + count: int = 1, + device: str = "cpu", +) -> MockArticulationView: + """Create a mock articulation view configured for a quadruped robot. + + Configuration: + - 12 DOFs (3 per leg x 4 legs: hip, thigh, calf) + - 13 links (base + 3 per leg) + - Floating base + + Args: + count: Number of articulation instances. + device: Device for tensor allocation. + + Returns: + A MockArticulationView configured for quadruped. + """ + dof_names = [ + "FL_hip_joint", + "FL_thigh_joint", + "FL_calf_joint", + "FR_hip_joint", + "FR_thigh_joint", + "FR_calf_joint", + "RL_hip_joint", + "RL_thigh_joint", + "RL_calf_joint", + "RR_hip_joint", + "RR_thigh_joint", + "RR_calf_joint", + ] + link_names = [ + "base", + "FL_hip", + "FL_thigh", + "FL_calf", + "FR_hip", + "FR_thigh", + "FR_calf", + "RL_hip", + "RL_thigh", + "RL_calf", + "RR_hip", + "RR_thigh", + "RR_calf", + ] + return MockArticulationView( + count=count, + num_dofs=12, + num_links=13, + dof_names=dof_names, + link_names=link_names, + fixed_base=False, + device=device, + ) + + +def create_mock_humanoid_view( + count: int = 1, + device: str = "cpu", +) -> MockArticulationView: + """Create a mock articulation view configured for a humanoid robot. + + Configuration: + - 21 DOFs (typical humanoid configuration) + - 22 links + - Floating base + + Args: + count: Number of articulation instances. + device: Device for tensor allocation. + + Returns: + A MockArticulationView configured for humanoid. + """ + dof_names = [ + # Torso + "torso_joint", + # Left arm + "left_shoulder_pitch", + "left_shoulder_roll", + "left_shoulder_yaw", + "left_elbow", + # Right arm + "right_shoulder_pitch", + "right_shoulder_roll", + "right_shoulder_yaw", + "right_elbow", + # Left leg + "left_hip_yaw", + "left_hip_roll", + "left_hip_pitch", + "left_knee", + "left_ankle_pitch", + "left_ankle_roll", + # Right leg + "right_hip_yaw", + "right_hip_roll", + "right_hip_pitch", + "right_knee", + "right_ankle_pitch", + "right_ankle_roll", + ] + link_names = [ + "pelvis", + "torso", + # Left arm + "left_shoulder", + "left_upper_arm", + "left_lower_arm", + "left_hand", + # Right arm + "right_shoulder", + "right_upper_arm", + "right_lower_arm", + "right_hand", + # Left leg + "left_hip", + "left_upper_leg", + "left_lower_leg", + "left_ankle", + "left_foot", + # Right leg + "right_hip", + "right_upper_leg", + "right_lower_leg", + "right_ankle", + "right_foot", + # Head + "neck", + "head", + ] + return MockArticulationView( + count=count, + num_dofs=21, + num_links=22, + dof_names=dof_names, + link_names=link_names, + fixed_base=False, + device=device, + ) diff --git a/source/isaaclab_physx/isaaclab_physx/test/mock_interfaces/utils/__init__.py b/source/isaaclab_physx/isaaclab_physx/test/mock_interfaces/utils/__init__.py new file mode 100644 index 00000000000..d945acdf8b1 --- /dev/null +++ b/source/isaaclab_physx/isaaclab_physx/test/mock_interfaces/utils/__init__.py @@ -0,0 +1,26 @@ +# 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 + +"""Utilities for mock PhysX interfaces.""" + +from .mock_shared_metatype import MockSharedMetatype +from .patching import ( + mock_articulation_view, + mock_rigid_body_view, + mock_rigid_contact_view, + patch_articulation_view, + patch_rigid_body_view, + patch_rigid_contact_view, +) + +__all__ = [ + "MockSharedMetatype", + "patch_rigid_body_view", + "patch_articulation_view", + "patch_rigid_contact_view", + "mock_rigid_body_view", + "mock_articulation_view", + "mock_rigid_contact_view", +] diff --git a/source/isaaclab_physx/isaaclab_physx/test/mock_interfaces/utils/mock_shared_metatype.py b/source/isaaclab_physx/isaaclab_physx/test/mock_interfaces/utils/mock_shared_metatype.py new file mode 100644 index 00000000000..c7ceb4d93f7 --- /dev/null +++ b/source/isaaclab_physx/isaaclab_physx/test/mock_interfaces/utils/mock_shared_metatype.py @@ -0,0 +1,64 @@ +# 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 + +"""Mock implementation of PhysX shared metatype.""" + +from __future__ import annotations + + +class MockSharedMetatype: + """Mock implementation of the shared metatype for articulation views. + + The shared metatype contains metadata about the articulation structure that is + shared across all instances, such as DOF count, link count, and names. + """ + + def __init__( + self, + dof_count: int = 1, + link_count: int = 2, + dof_names: list[str] | None = None, + link_names: list[str] | None = None, + fixed_base: bool = False, + ): + """Initialize the mock shared metatype. + + Args: + dof_count: Number of degrees of freedom (joints). + link_count: Number of links (bodies). + dof_names: Names of the DOFs. Defaults to ["dof_0", "dof_1", ...]. + link_names: Names of the links. Defaults to ["link_0", "link_1", ...]. + fixed_base: Whether the articulation has a fixed base. + """ + self._dof_count = dof_count + self._link_count = link_count + self._dof_names = dof_names or [f"dof_{i}" for i in range(dof_count)] + self._link_names = link_names or [f"link_{i}" for i in range(link_count)] + self._fixed_base = fixed_base + + @property + def dof_count(self) -> int: + """Number of degrees of freedom.""" + return self._dof_count + + @property + def link_count(self) -> int: + """Number of links.""" + return self._link_count + + @property + def dof_names(self) -> list[str]: + """Names of the degrees of freedom.""" + return self._dof_names + + @property + def link_names(self) -> list[str]: + """Names of the links.""" + return self._link_names + + @property + def fixed_base(self) -> bool: + """Whether the articulation has a fixed base.""" + return self._fixed_base diff --git a/source/isaaclab_physx/isaaclab_physx/test/mock_interfaces/utils/patching.py b/source/isaaclab_physx/isaaclab_physx/test/mock_interfaces/utils/patching.py new file mode 100644 index 00000000000..288eb28bf5f --- /dev/null +++ b/source/isaaclab_physx/isaaclab_physx/test/mock_interfaces/utils/patching.py @@ -0,0 +1,285 @@ +# 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 + +"""Utilities for patching PhysX views with mock implementations.""" + +from __future__ import annotations + +import functools +from collections.abc import Callable, Generator +from contextlib import contextmanager +from typing import Any, TypeVar +from unittest.mock import patch + +F = TypeVar("F", bound=Callable[..., Any]) + + +@contextmanager +def patch_rigid_body_view( + target: str, + count: int = 1, + prim_paths: list[str] | None = None, + device: str = "cpu", +) -> Generator[Any, None, None]: + """Context manager for patching physx.RigidBodyView with a mock. + + Args: + target: The target to patch (e.g., "my_module.physx.RigidBodyView"). + count: Number of rigid body instances. + prim_paths: USD prim paths for each instance. + device: Device for tensor allocation. + + Yields: + The mock class that will be used when the target is instantiated. + + Example: + >>> with patch_rigid_body_view("my_module.view", count=4) as mock: + ... view = my_function() + ... view.set_mock_transforms(torch.randn(4, 7)) + """ + from ..views import MockRigidBodyView + + def create_mock(*args, **kwargs): + return MockRigidBodyView( + count=kwargs.get("count", count), + prim_paths=kwargs.get("prim_paths", prim_paths), + device=kwargs.get("device", device), + ) + + with patch(target, side_effect=create_mock) as mock_class: + yield mock_class + + +@contextmanager +def patch_articulation_view( + target: str, + count: int = 1, + num_dofs: int = 1, + num_links: int = 2, + dof_names: list[str] | None = None, + link_names: list[str] | None = None, + fixed_base: bool = False, + prim_paths: list[str] | None = None, + device: str = "cpu", +) -> Generator[Any, None, None]: + """Context manager for patching physx.ArticulationView with a mock. + + Args: + target: The target to patch (e.g., "my_module.physx.ArticulationView"). + count: Number of articulation instances. + num_dofs: Number of degrees of freedom (joints). + num_links: Number of links (bodies). + dof_names: Names of the DOFs. + link_names: Names of the links. + fixed_base: Whether the articulation has a fixed base. + prim_paths: USD prim paths for each instance. + device: Device for tensor allocation. + + Yields: + The mock class that will be used when the target is instantiated. + + Example: + >>> with patch_articulation_view("my_module.view", num_dofs=12) as mock: + ... view = my_function() + ... positions = view.get_dof_positions() + """ + from ..views import MockArticulationView + + def create_mock(*args, **kwargs): + return MockArticulationView( + count=kwargs.get("count", count), + num_dofs=kwargs.get("num_dofs", num_dofs), + num_links=kwargs.get("num_links", num_links), + dof_names=kwargs.get("dof_names", dof_names), + link_names=kwargs.get("link_names", link_names), + fixed_base=kwargs.get("fixed_base", fixed_base), + prim_paths=kwargs.get("prim_paths", prim_paths), + device=kwargs.get("device", device), + ) + + with patch(target, side_effect=create_mock) as mock_class: + yield mock_class + + +@contextmanager +def patch_rigid_contact_view( + target: str, + count: int = 1, + num_bodies: int = 1, + filter_count: int = 0, + max_contact_data_count: int = 16, + device: str = "cpu", +) -> Generator[Any, None, None]: + """Context manager for patching physx.RigidContactView with a mock. + + Args: + target: The target to patch (e.g., "my_module.physx.RigidContactView"). + count: Number of instances. + num_bodies: Number of bodies per instance. + filter_count: Number of filter bodies for contact filtering. + max_contact_data_count: Maximum number of contact data points. + device: Device for tensor allocation. + + Yields: + The mock class that will be used when the target is instantiated. + + Example: + >>> with patch_rigid_contact_view("my_module.view", num_bodies=4) as mock: + ... view = my_function() + ... forces = view.get_net_contact_forces(0.01) + """ + from ..views import MockRigidContactView + + def create_mock(*args, **kwargs): + return MockRigidContactView( + count=kwargs.get("count", count), + num_bodies=kwargs.get("num_bodies", num_bodies), + filter_count=kwargs.get("filter_count", filter_count), + max_contact_data_count=kwargs.get("max_contact_data_count", max_contact_data_count), + device=kwargs.get("device", device), + ) + + with patch(target, side_effect=create_mock) as mock_class: + yield mock_class + + +# -- Decorators -- + + +def mock_rigid_body_view( + count: int = 1, + prim_paths: list[str] | None = None, + device: str = "cpu", +) -> Callable[[F], F]: + """Decorator for injecting MockRigidBodyView into test function. + + The mock view is passed as the first argument to the decorated function. + + Args: + count: Number of rigid body instances. + prim_paths: USD prim paths for each instance. + device: Device for tensor allocation. + + Returns: + A decorator function. + + Example: + >>> @mock_rigid_body_view(count=4) + ... def test_my_function(mock_view): + ... transforms = mock_view.get_transforms() + ... assert transforms.shape == (4, 7) + """ + from ..views import MockRigidBodyView + + def decorator(func: F) -> F: + @functools.wraps(func) + def wrapper(*args, **kwargs): + mock = MockRigidBodyView(count=count, prim_paths=prim_paths, device=device) + return func(mock, *args, **kwargs) + + return wrapper # type: ignore[return-value] + + return decorator + + +def mock_articulation_view( + count: int = 1, + num_dofs: int = 1, + num_links: int = 2, + dof_names: list[str] | None = None, + link_names: list[str] | None = None, + fixed_base: bool = False, + prim_paths: list[str] | None = None, + device: str = "cpu", +) -> Callable[[F], F]: + """Decorator for injecting MockArticulationView into test function. + + The mock view is passed as the first argument to the decorated function. + + Args: + count: Number of articulation instances. + num_dofs: Number of degrees of freedom (joints). + num_links: Number of links (bodies). + dof_names: Names of the DOFs. + link_names: Names of the links. + fixed_base: Whether the articulation has a fixed base. + prim_paths: USD prim paths for each instance. + device: Device for tensor allocation. + + Returns: + A decorator function. + + Example: + >>> @mock_articulation_view(count=4, num_dofs=12, num_links=13) + ... def test_my_function(mock_view): + ... positions = mock_view.get_dof_positions() + ... assert positions.shape == (4, 12) + """ + from ..views import MockArticulationView + + def decorator(func: F) -> F: + @functools.wraps(func) + def wrapper(*args, **kwargs): + mock = MockArticulationView( + count=count, + num_dofs=num_dofs, + num_links=num_links, + dof_names=dof_names, + link_names=link_names, + fixed_base=fixed_base, + prim_paths=prim_paths, + device=device, + ) + return func(mock, *args, **kwargs) + + return wrapper # type: ignore[return-value] + + return decorator + + +def mock_rigid_contact_view( + count: int = 1, + num_bodies: int = 1, + filter_count: int = 0, + max_contact_data_count: int = 16, + device: str = "cpu", +) -> Callable[[F], F]: + """Decorator for injecting MockRigidContactView into test function. + + The mock view is passed as the first argument to the decorated function. + + Args: + count: Number of instances. + num_bodies: Number of bodies per instance. + filter_count: Number of filter bodies for contact filtering. + max_contact_data_count: Maximum number of contact data points. + device: Device for tensor allocation. + + Returns: + A decorator function. + + Example: + >>> @mock_rigid_contact_view(count=4, num_bodies=4, filter_count=2) + ... def test_my_function(mock_view): + ... forces = mock_view.get_net_contact_forces(0.01) + ... assert forces.shape == (16, 3) + """ + from ..views import MockRigidContactView + + def decorator(func: F) -> F: + @functools.wraps(func) + def wrapper(*args, **kwargs): + mock = MockRigidContactView( + count=count, + num_bodies=num_bodies, + filter_count=filter_count, + max_contact_data_count=max_contact_data_count, + device=device, + ) + return func(mock, *args, **kwargs) + + return wrapper # type: ignore[return-value] + + return decorator diff --git a/source/isaaclab_physx/isaaclab_physx/test/mock_interfaces/views/__init__.py b/source/isaaclab_physx/isaaclab_physx/test/mock_interfaces/views/__init__.py new file mode 100644 index 00000000000..f4699c04592 --- /dev/null +++ b/source/isaaclab_physx/isaaclab_physx/test/mock_interfaces/views/__init__.py @@ -0,0 +1,12 @@ +# 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 + +"""Mock PhysX TensorAPI views.""" + +from .mock_articulation_view import MockArticulationView +from .mock_rigid_body_view import MockRigidBodyView +from .mock_rigid_contact_view import MockRigidContactView + +__all__ = ["MockRigidBodyView", "MockArticulationView", "MockRigidContactView"] diff --git a/source/isaaclab_physx/isaaclab_physx/test/mock_interfaces/views/mock_articulation_view.py b/source/isaaclab_physx/isaaclab_physx/test/mock_interfaces/views/mock_articulation_view.py new file mode 100644 index 00000000000..52827ac3c53 --- /dev/null +++ b/source/isaaclab_physx/isaaclab_physx/test/mock_interfaces/views/mock_articulation_view.py @@ -0,0 +1,917 @@ +# 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 + +"""Mock implementation of PhysX ArticulationView.""" + +from __future__ import annotations + +import torch + +from ..utils.mock_shared_metatype import MockSharedMetatype + + +class MockArticulationView: + """Mock implementation of physx.ArticulationView for unit testing. + + This class mimics the interface of the PhysX TensorAPI ArticulationView, + allowing tests to run without Isaac Sim or GPU simulation. + + Data Shapes: + - root_transforms: (N, 7) - [pos(3), quat_xyzw(4)] + - root_velocities: (N, 6) - [lin_vel(3), ang_vel(3)] + - link_transforms: (N, L, 7) - per-link poses + - link_velocities: (N, L, 6) - per-link velocities + - dof_positions: (N, J) - joint positions + - dof_velocities: (N, J) - joint velocities + - dof_limits: (N, J, 2) - [lower, upper] limits + - dof_stiffnesses: (N, J) - joint stiffnesses + - dof_dampings: (N, J) - joint dampings + - dof_max_forces: (N, J) - maximum joint forces + - dof_max_velocities: (N, J) - maximum joint velocities + - masses: (N, L) - per-link masses + - coms: (N, L, 7) - per-link centers of mass + - inertias: (N, L, 3, 3) - per-link inertia tensors + + Where N = count, L = num_links, J = num_dofs + """ + + def __init__( + self, + count: int = 1, + num_dofs: int = 1, + num_links: int = 2, + dof_names: list[str] | None = None, + link_names: list[str] | None = None, + fixed_base: bool = False, + prim_paths: list[str] | None = None, + device: str = "cpu", + ): + """Initialize the mock articulation view. + + Args: + count: Number of articulation instances. + num_dofs: Number of degrees of freedom (joints). + num_links: Number of links (bodies). + dof_names: Names of the DOFs. Defaults to auto-generated names. + link_names: Names of the links. Defaults to auto-generated names. + fixed_base: Whether the articulation has a fixed base. + prim_paths: USD prim paths for each instance. + device: Device for tensor allocation ("cpu" or "cuda"). + """ + self._count = count + self._num_dofs = num_dofs + self._num_links = num_links + self._device = device + self._prim_paths = prim_paths or [f"/World/Articulation_{i}" for i in range(count)] + + # Create shared metatype + self._shared_metatype = MockSharedMetatype( + dof_count=num_dofs, + link_count=num_links, + dof_names=dof_names, + link_names=link_names, + fixed_base=fixed_base, + ) + + # Tendon properties (fixed values for mock) + self._max_fixed_tendons = 0 + self._max_spatial_tendons = 0 + + # Internal state (lazily initialized) + self._root_transforms: torch.Tensor | None = None + self._root_velocities: torch.Tensor | None = None + self._link_transforms: torch.Tensor | None = None + self._link_velocities: torch.Tensor | None = None + self._link_accelerations: torch.Tensor | None = None + self._link_incoming_joint_force: torch.Tensor | None = None + self._dof_positions: torch.Tensor | None = None + self._dof_velocities: torch.Tensor | None = None + self._dof_projected_joint_forces: torch.Tensor | None = None + self._dof_limits: torch.Tensor | None = None + self._dof_stiffnesses: torch.Tensor | None = None + self._dof_dampings: torch.Tensor | None = None + self._dof_max_forces: torch.Tensor | None = None + self._dof_max_velocities: torch.Tensor | None = None + self._dof_armatures: torch.Tensor | None = None + self._dof_friction_coefficients: torch.Tensor | None = None + self._dof_friction_properties: torch.Tensor | None = None + self._masses: torch.Tensor | None = None + self._coms: torch.Tensor | None = None + self._inertias: torch.Tensor | None = None + + # -- Helper Methods -- + + def _check_cpu_tensor(self, tensor: torch.Tensor, name: str) -> None: + """Check that tensor is on CPU, raise RuntimeError if on GPU. + + This mimics PhysX behavior where joint/body properties must be on CPU. + """ + if tensor.is_cuda: + raise RuntimeError( + f"Expected CPU tensor for {name}, but got tensor on {tensor.device}. " + "Joint and body properties must be set with CPU tensors." + ) + + # -- Properties -- + + @property + def count(self) -> int: + """Number of articulation instances.""" + return self._count + + @property + def shared_metatype(self) -> MockSharedMetatype: + """Shared metatype containing articulation structure metadata.""" + return self._shared_metatype + + @property + def max_fixed_tendons(self) -> int: + """Maximum number of fixed tendons.""" + return self._max_fixed_tendons + + @property + def max_spatial_tendons(self) -> int: + """Maximum number of spatial tendons.""" + return self._max_spatial_tendons + + @property + def prim_paths(self) -> list[str]: + """USD prim paths for each instance.""" + return self._prim_paths + + # -- Root Getters -- + + def get_root_transforms(self) -> torch.Tensor: + """Get world transforms of root links. + + Returns: + Tensor of shape (N, 7) with [pos(3), quat_xyzw(4)]. + """ + if self._root_transforms is None: + self._root_transforms = torch.zeros(self._count, 7, device=self._device) + self._root_transforms[:, 6] = 1.0 # w=1 for identity quaternion + return self._root_transforms.clone() + + def get_root_velocities(self) -> torch.Tensor: + """Get velocities of root links. + + Returns: + Tensor of shape (N, 6) with [lin_vel(3), ang_vel(3)]. + """ + if self._root_velocities is None: + self._root_velocities = torch.zeros(self._count, 6, device=self._device) + return self._root_velocities.clone() + + # -- Link Getters -- + + def get_link_transforms(self) -> torch.Tensor: + """Get world transforms of all links. + + Returns: + Tensor of shape (N, L, 7) with [pos(3), quat_xyzw(4)] per link. + """ + if self._link_transforms is None: + self._link_transforms = torch.zeros(self._count, self._num_links, 7, device=self._device) + self._link_transforms[:, :, 6] = 1.0 # w=1 for identity quaternion + return self._link_transforms.clone() + + def get_link_velocities(self) -> torch.Tensor: + """Get velocities of all links. + + Returns: + Tensor of shape (N, L, 6) with [lin_vel(3), ang_vel(3)] per link. + """ + if self._link_velocities is None: + self._link_velocities = torch.zeros(self._count, self._num_links, 6, device=self._device) + return self._link_velocities.clone() + + # -- DOF Getters -- + + def get_dof_positions(self) -> torch.Tensor: + """Get positions of all DOFs. + + Returns: + Tensor of shape (N, J) with joint positions. + """ + if self._dof_positions is None: + self._dof_positions = torch.zeros(self._count, self._num_dofs, device=self._device) + return self._dof_positions.clone() + + def get_dof_velocities(self) -> torch.Tensor: + """Get velocities of all DOFs. + + Returns: + Tensor of shape (N, J) with joint velocities. + """ + if self._dof_velocities is None: + self._dof_velocities = torch.zeros(self._count, self._num_dofs, device=self._device) + return self._dof_velocities.clone() + + def get_dof_projected_joint_forces(self) -> torch.Tensor: + """Get projected joint forces of all DOFs. + + Returns: + Tensor of shape (N, J) with projected joint forces. + """ + if self._dof_projected_joint_forces is None: + self._dof_projected_joint_forces = torch.zeros(self._count, self._num_dofs, device=self._device) + return self._dof_projected_joint_forces.clone() + + def get_dof_limits(self) -> torch.Tensor: + """Get position limits of all DOFs. + + Returns: + Tensor of shape (N, J, 2) with [lower, upper] limits. Always on CPU. + """ + if self._dof_limits is None: + # Default: no limits (infinite) - stored on CPU + self._dof_limits = torch.zeros(self._count, self._num_dofs, 2, device="cpu") + self._dof_limits[:, :, 0] = float("-inf") # lower limit + self._dof_limits[:, :, 1] = float("inf") # upper limit + return self._dof_limits.clone() + + def get_dof_stiffnesses(self) -> torch.Tensor: + """Get stiffnesses of all DOFs. + + Returns: + Tensor of shape (N, J) with joint stiffnesses. Always on CPU. + """ + if self._dof_stiffnesses is None: + self._dof_stiffnesses = torch.zeros(self._count, self._num_dofs, device="cpu") + return self._dof_stiffnesses.clone() + + def get_dof_dampings(self) -> torch.Tensor: + """Get dampings of all DOFs. + + Returns: + Tensor of shape (N, J) with joint dampings. Always on CPU. + """ + if self._dof_dampings is None: + self._dof_dampings = torch.zeros(self._count, self._num_dofs, device="cpu") + return self._dof_dampings.clone() + + def get_dof_max_forces(self) -> torch.Tensor: + """Get maximum forces of all DOFs. + + Returns: + Tensor of shape (N, J) with maximum joint forces. Always on CPU. + """ + if self._dof_max_forces is None: + # Default: infinite max force - stored on CPU + self._dof_max_forces = torch.full((self._count, self._num_dofs), float("inf"), device="cpu") + return self._dof_max_forces.clone() + + def get_dof_max_velocities(self) -> torch.Tensor: + """Get maximum velocities of all DOFs. + + Returns: + Tensor of shape (N, J) with maximum joint velocities. Always on CPU. + """ + if self._dof_max_velocities is None: + # Default: infinite max velocity - stored on CPU + self._dof_max_velocities = torch.full((self._count, self._num_dofs), float("inf"), device="cpu") + return self._dof_max_velocities.clone() + + def get_dof_armatures(self) -> torch.Tensor: + """Get armatures of all DOFs. + + Returns: + Tensor of shape (N, J) with joint armatures. Always on CPU. + """ + if self._dof_armatures is None: + self._dof_armatures = torch.zeros(self._count, self._num_dofs, device="cpu") + return self._dof_armatures.clone() + + def get_dof_friction_coefficients(self) -> torch.Tensor: + """Get friction coefficients of all DOFs. + + Returns: + Tensor of shape (N, J) with joint friction coefficients. Always on CPU. + """ + if self._dof_friction_coefficients is None: + self._dof_friction_coefficients = torch.zeros(self._count, self._num_dofs, device="cpu") + return self._dof_friction_coefficients.clone() + + # -- Mass Property Getters -- + + def get_masses(self) -> torch.Tensor: + """Get masses of all links. + + Returns: + Tensor of shape (N, L) with link masses. Always on CPU. + """ + if self._masses is None: + self._masses = torch.ones(self._count, self._num_links, device="cpu") + return self._masses.clone() + + def get_coms(self) -> torch.Tensor: + """Get centers of mass of all links. + + Returns: + Tensor of shape (N, L, 7) with [pos(3), quat_xyzw(4)] per link. Always on CPU. + """ + if self._coms is None: + self._coms = torch.zeros(self._count, self._num_links, 7, device="cpu") + self._coms[:, :, 6] = 1.0 # w=1 for identity quaternion + return self._coms.clone() + + def get_inertias(self) -> torch.Tensor: + """Get inertia tensors of all links. + + Returns: + Tensor of shape (N, L, 9) with flattened 3x3 inertia matrices per link (row-major). Always on CPU. + """ + if self._inertias is None: + # Default: identity inertia - flattened [1,0,0,0,1,0,0,0,1] - stored on CPU + self._inertias = torch.zeros(self._count, self._num_links, 9, device="cpu") + self._inertias[:, :, 0] = 1.0 # [0,0] + self._inertias[:, :, 4] = 1.0 # [1,1] + self._inertias[:, :, 8] = 1.0 # [2,2] + return self._inertias.clone() + + # -- Root Setters -- + + def set_root_transforms( + self, + transforms: torch.Tensor, + indices: torch.Tensor | None = None, + ) -> None: + """Set world transforms of root links. + + Args: + transforms: Tensor of shape (N, 7) or (len(indices), 7). + indices: Optional indices of articulations to update. + """ + transforms = transforms.to(self._device) + if self._root_transforms is None: + self._root_transforms = torch.zeros(self._count, 7, device=self._device) + self._root_transforms[:, 6] = 1.0 + if indices is not None: + self._root_transforms[indices] = transforms + else: + self._root_transforms = transforms + + def set_root_velocities( + self, + velocities: torch.Tensor, + indices: torch.Tensor | None = None, + ) -> None: + """Set velocities of root links. + + Args: + velocities: Tensor of shape (N, 6) or (len(indices), 6). + indices: Optional indices of articulations to update. + """ + velocities = velocities.to(self._device) + if self._root_velocities is None: + self._root_velocities = torch.zeros(self._count, 6, device=self._device) + if indices is not None: + self._root_velocities[indices] = velocities + else: + self._root_velocities = velocities + + # -- DOF Setters -- + + def set_dof_positions( + self, + positions: torch.Tensor, + indices: torch.Tensor | None = None, + ) -> None: + """Set positions of all DOFs. + + Args: + positions: Tensor of shape (N, J) or (len(indices), J). + indices: Optional indices of articulations to update. + """ + positions = positions.to(self._device) + if self._dof_positions is None: + self._dof_positions = torch.zeros(self._count, self._num_dofs, device=self._device) + if indices is not None: + self._dof_positions[indices] = positions + else: + self._dof_positions = positions + + def set_dof_velocities( + self, + velocities: torch.Tensor, + indices: torch.Tensor | None = None, + ) -> None: + """Set velocities of all DOFs. + + Args: + velocities: Tensor of shape (N, J) or (len(indices), J). + indices: Optional indices of articulations to update. + """ + velocities = velocities.to(self._device) + if self._dof_velocities is None: + self._dof_velocities = torch.zeros(self._count, self._num_dofs, device=self._device) + if indices is not None: + self._dof_velocities[indices] = velocities + else: + self._dof_velocities = velocities + + def set_dof_position_targets( + self, + targets: torch.Tensor, + indices: torch.Tensor | None = None, + ) -> None: + """Set position targets for all DOFs (no-op in mock). + + Args: + targets: Tensor of shape (N, J) or (len(indices), J). + indices: Optional indices of articulations to update. + """ + pass # No-op for mock + + def set_dof_velocity_targets( + self, + targets: torch.Tensor, + indices: torch.Tensor | None = None, + ) -> None: + """Set velocity targets for all DOFs (no-op in mock). + + Args: + targets: Tensor of shape (N, J) or (len(indices), J). + indices: Optional indices of articulations to update. + """ + pass # No-op for mock + + def set_dof_actuation_forces( + self, + forces: torch.Tensor, + indices: torch.Tensor | None = None, + ) -> None: + """Set actuation forces for all DOFs (no-op in mock). + + Args: + forces: Tensor of shape (N, J) or (len(indices), J). + indices: Optional indices of articulations to update. + """ + pass # No-op for mock + + def set_dof_limits( + self, + limits: torch.Tensor, + indices: torch.Tensor | None = None, + ) -> None: + """Set position limits of all DOFs. + + Args: + limits: Tensor of shape (N, J, 2) with [lower, upper] limits. Must be on CPU. + indices: Optional indices of articulations to update. + + Raises: + RuntimeError: If limits tensor is on GPU. + """ + self._check_cpu_tensor(limits, "dof_limits") + if self._dof_limits is None: + self._dof_limits = torch.zeros(self._count, self._num_dofs, 2, device="cpu") + self._dof_limits[:, :, 0] = float("-inf") + self._dof_limits[:, :, 1] = float("inf") + if indices is not None: + self._dof_limits[indices] = limits + else: + self._dof_limits = limits + + def set_dof_stiffnesses( + self, + stiffnesses: torch.Tensor, + indices: torch.Tensor | None = None, + ) -> None: + """Set stiffnesses of all DOFs. + + Args: + stiffnesses: Tensor of shape (N, J). Must be on CPU. + indices: Optional indices of articulations to update. + + Raises: + RuntimeError: If stiffnesses tensor is on GPU. + """ + self._check_cpu_tensor(stiffnesses, "dof_stiffnesses") + if self._dof_stiffnesses is None: + self._dof_stiffnesses = torch.zeros(self._count, self._num_dofs, device="cpu") + if indices is not None: + self._dof_stiffnesses[indices] = stiffnesses + else: + self._dof_stiffnesses = stiffnesses + + def set_dof_dampings( + self, + dampings: torch.Tensor, + indices: torch.Tensor | None = None, + ) -> None: + """Set dampings of all DOFs. + + Args: + dampings: Tensor of shape (N, J). Must be on CPU. + indices: Optional indices of articulations to update. + + Raises: + RuntimeError: If dampings tensor is on GPU. + """ + self._check_cpu_tensor(dampings, "dof_dampings") + if self._dof_dampings is None: + self._dof_dampings = torch.zeros(self._count, self._num_dofs, device="cpu") + if indices is not None: + self._dof_dampings[indices] = dampings + else: + self._dof_dampings = dampings + + def set_dof_max_forces( + self, + max_forces: torch.Tensor, + indices: torch.Tensor | None = None, + ) -> None: + """Set maximum forces of all DOFs. + + Args: + max_forces: Tensor of shape (N, J). Must be on CPU. + indices: Optional indices of articulations to update. + + Raises: + RuntimeError: If max_forces tensor is on GPU. + """ + self._check_cpu_tensor(max_forces, "dof_max_forces") + if self._dof_max_forces is None: + self._dof_max_forces = torch.full((self._count, self._num_dofs), float("inf"), device="cpu") + if indices is not None: + self._dof_max_forces[indices] = max_forces + else: + self._dof_max_forces = max_forces + + def set_dof_max_velocities( + self, + max_velocities: torch.Tensor, + indices: torch.Tensor | None = None, + ) -> None: + """Set maximum velocities of all DOFs. + + Args: + max_velocities: Tensor of shape (N, J). Must be on CPU. + indices: Optional indices of articulations to update. + + Raises: + RuntimeError: If max_velocities tensor is on GPU. + """ + self._check_cpu_tensor(max_velocities, "dof_max_velocities") + if self._dof_max_velocities is None: + self._dof_max_velocities = torch.full((self._count, self._num_dofs), float("inf"), device="cpu") + if indices is not None: + self._dof_max_velocities[indices] = max_velocities + else: + self._dof_max_velocities = max_velocities + + def set_dof_armatures( + self, + armatures: torch.Tensor, + indices: torch.Tensor | None = None, + ) -> None: + """Set armatures of all DOFs. + + Args: + armatures: Tensor of shape (N, J). Must be on CPU. + indices: Optional indices of articulations to update. + + Raises: + RuntimeError: If armatures tensor is on GPU. + """ + self._check_cpu_tensor(armatures, "dof_armatures") + if self._dof_armatures is None: + self._dof_armatures = torch.zeros(self._count, self._num_dofs, device="cpu") + if indices is not None: + self._dof_armatures[indices] = armatures + else: + self._dof_armatures = armatures + + def set_dof_friction_coefficients( + self, + friction_coefficients: torch.Tensor, + indices: torch.Tensor | None = None, + ) -> None: + """Set friction coefficients of all DOFs. + + Args: + friction_coefficients: Tensor of shape (N, J). Must be on CPU. + indices: Optional indices of articulations to update. + + Raises: + RuntimeError: If friction_coefficients tensor is on GPU. + """ + self._check_cpu_tensor(friction_coefficients, "dof_friction_coefficients") + if self._dof_friction_coefficients is None: + self._dof_friction_coefficients = torch.zeros(self._count, self._num_dofs, device="cpu") + if indices is not None: + self._dof_friction_coefficients[indices] = friction_coefficients + else: + self._dof_friction_coefficients = friction_coefficients + + # -- Mass Property Setters -- + + def set_masses( + self, + masses: torch.Tensor, + indices: torch.Tensor | None = None, + ) -> None: + """Set masses of all links. + + Args: + masses: Tensor of shape (N, L). Must be on CPU. + indices: Optional indices of articulations to update. + + Raises: + RuntimeError: If masses tensor is on GPU. + """ + self._check_cpu_tensor(masses, "masses") + if self._masses is None: + self._masses = torch.ones(self._count, self._num_links, device="cpu") + if indices is not None: + self._masses[indices] = masses + else: + self._masses = masses + + def set_coms( + self, + coms: torch.Tensor, + indices: torch.Tensor | None = None, + ) -> None: + """Set centers of mass of all links. + + Args: + coms: Tensor of shape (N, L, 7). Must be on CPU. + indices: Optional indices of articulations to update. + + Raises: + RuntimeError: If coms tensor is on GPU. + """ + self._check_cpu_tensor(coms, "coms") + if self._coms is None: + self._coms = torch.zeros(self._count, self._num_links, 7, device="cpu") + self._coms[:, :, 6] = 1.0 + if indices is not None: + self._coms[indices] = coms + else: + self._coms = coms + + def set_inertias( + self, + inertias: torch.Tensor, + indices: torch.Tensor | None = None, + ) -> None: + """Set inertia tensors of all links. + + Args: + inertias: Tensor of shape (N, L, 9) - flattened 3x3 matrices. Must be on CPU. + indices: Optional indices of articulations to update. + + Raises: + RuntimeError: If inertias tensor is on GPU. + """ + self._check_cpu_tensor(inertias, "inertias") + if self._inertias is None: + self._inertias = torch.zeros(self._count, self._num_links, 9, device="cpu") + self._inertias[:, :, 0] = 1.0 + self._inertias[:, :, 4] = 1.0 + self._inertias[:, :, 8] = 1.0 + if indices is not None: + self._inertias[indices] = inertias + else: + self._inertias = inertias + + # -- Mock setters (direct test data injection) -- + + def set_mock_root_transforms(self, transforms: torch.Tensor) -> None: + """Set mock root transform data directly for testing. + + Args: + transforms: Tensor of shape (N, 7). + """ + self._root_transforms = transforms.to(self._device) + + def set_mock_root_velocities(self, velocities: torch.Tensor) -> None: + """Set mock root velocity data directly for testing. + + Args: + velocities: Tensor of shape (N, 6). + """ + self._root_velocities = velocities.to(self._device) + + def set_mock_link_transforms(self, transforms: torch.Tensor) -> None: + """Set mock link transform data directly for testing. + + Args: + transforms: Tensor of shape (N, L, 7). + """ + self._link_transforms = transforms.to(self._device) + + def set_mock_link_velocities(self, velocities: torch.Tensor) -> None: + """Set mock link velocity data directly for testing. + + Args: + velocities: Tensor of shape (N, L, 6). + """ + self._link_velocities = velocities.to(self._device) + + def set_mock_dof_positions(self, positions: torch.Tensor) -> None: + """Set mock DOF position data directly for testing. + + Args: + positions: Tensor of shape (N, J). + """ + self._dof_positions = positions.to(self._device) + + def set_mock_dof_velocities(self, velocities: torch.Tensor) -> None: + """Set mock DOF velocity data directly for testing. + + Args: + velocities: Tensor of shape (N, J). + """ + self._dof_velocities = velocities.to(self._device) + + def set_mock_dof_projected_joint_forces(self, forces: torch.Tensor) -> None: + """Set mock projected joint force data directly for testing. + + Args: + forces: Tensor of shape (N, J). + """ + self._dof_projected_joint_forces = forces.to(self._device) + + def set_mock_dof_limits(self, limits: torch.Tensor) -> None: + """Set mock DOF limit data directly for testing. + + Args: + limits: Tensor of shape (N, J, 2). + """ + self._dof_limits = limits.to(self._device) + + def set_mock_dof_stiffnesses(self, stiffnesses: torch.Tensor) -> None: + """Set mock DOF stiffness data directly for testing. + + Args: + stiffnesses: Tensor of shape (N, J). + """ + self._dof_stiffnesses = stiffnesses.to(self._device) + + def set_mock_dof_dampings(self, dampings: torch.Tensor) -> None: + """Set mock DOF damping data directly for testing. + + Args: + dampings: Tensor of shape (N, J). + """ + self._dof_dampings = dampings.to(self._device) + + def set_mock_dof_max_forces(self, max_forces: torch.Tensor) -> None: + """Set mock DOF max force data directly for testing. + + Args: + max_forces: Tensor of shape (N, J). + """ + self._dof_max_forces = max_forces.to(self._device) + + def set_mock_dof_max_velocities(self, max_velocities: torch.Tensor) -> None: + """Set mock DOF max velocity data directly for testing. + + Args: + max_velocities: Tensor of shape (N, J). + """ + self._dof_max_velocities = max_velocities.to(self._device) + + def set_mock_dof_armatures(self, armatures: torch.Tensor) -> None: + """Set mock DOF armature data directly for testing. + + Args: + armatures: Tensor of shape (N, J). + """ + self._dof_armatures = armatures.to(self._device) + + def set_mock_dof_friction_coefficients(self, friction_coefficients: torch.Tensor) -> None: + """Set mock DOF friction coefficient data directly for testing. + + Args: + friction_coefficients: Tensor of shape (N, J). + """ + self._dof_friction_coefficients = friction_coefficients.to(self._device) + + def set_mock_masses(self, masses: torch.Tensor) -> None: + """Set mock mass data directly for testing. + + Args: + masses: Tensor of shape (N, L). + """ + self._masses = masses.to(self._device) + + def set_mock_coms(self, coms: torch.Tensor) -> None: + """Set mock center of mass data directly for testing. + + Args: + coms: Tensor of shape (N, L, 7). + """ + self._coms = coms.to(self._device) + + def set_mock_inertias(self, inertias: torch.Tensor) -> None: + """Set mock inertia data directly for testing. + + Args: + inertias: Tensor of shape (N, L, 3, 3). + """ + self._inertias = inertias.to(self._device) + + # -- Additional mock state for extended properties -- + + def get_dof_friction_properties(self) -> torch.Tensor: + """Get friction properties of all DOFs. + + Returns: + Tensor of shape (N, J, 3) with [static_friction, dynamic_friction, viscous_friction]. Always on CPU. + """ + if self._dof_friction_properties is None: + self._dof_friction_properties = torch.zeros(self._count, self._num_dofs, 3, device="cpu") + return self._dof_friction_properties.clone() + + def get_link_accelerations(self) -> torch.Tensor: + """Get accelerations of all links. + + Returns: + Tensor of shape (N, L, 6) with [lin_acc(3), ang_acc(3)] per link. + """ + if self._link_accelerations is None: + self._link_accelerations = torch.zeros(self._count, self._num_links, 6, device=self._device) + return self._link_accelerations.clone() + + def get_link_incoming_joint_force(self) -> torch.Tensor: + """Get incoming joint forces for all links. + + Returns: + Tensor of shape (N, L, 6) with [force(3), torque(3)] per link. + """ + if self._link_incoming_joint_force is None: + self._link_incoming_joint_force = torch.zeros(self._count, self._num_links, 6, device=self._device) + return self._link_incoming_joint_force.clone() + + def set_mock_dof_friction_properties(self, friction_properties: torch.Tensor) -> None: + """Set mock DOF friction properties data directly for testing. + + Args: + friction_properties: Tensor of shape (N, J, 3). + """ + self._dof_friction_properties = friction_properties.to(self._device) + + def set_mock_link_accelerations(self, accelerations: torch.Tensor) -> None: + """Set mock link acceleration data directly for testing. + + Args: + accelerations: Tensor of shape (N, L, 6). + """ + self._link_accelerations = accelerations.to(self._device) + + def set_mock_link_incoming_joint_force(self, forces: torch.Tensor) -> None: + """Set mock link incoming joint force data directly for testing. + + Args: + forces: Tensor of shape (N, L, 6). + """ + self._link_incoming_joint_force = forces.to(self._device) + + def set_random_mock_data(self) -> None: + """Set all internal state to random values for benchmarking. + + This method initializes all mock data with random values, + useful for benchmarking where the actual values don't matter. + """ + # Root state + self._root_transforms = torch.randn(self._count, 7, device=self._device) + self._root_transforms[:, 3:7] = torch.nn.functional.normalize(self._root_transforms[:, 3:7], dim=-1) + self._root_velocities = torch.randn(self._count, 6, device=self._device) + + # Link state + self._link_transforms = torch.randn(self._count, self._num_links, 7, device=self._device) + self._link_transforms[:, :, 3:7] = torch.nn.functional.normalize(self._link_transforms[:, :, 3:7], dim=-1) + self._link_velocities = torch.randn(self._count, self._num_links, 6, device=self._device) + self._link_accelerations = torch.randn(self._count, self._num_links, 6, device=self._device) + self._link_incoming_joint_force = torch.randn(self._count, self._num_links, 6, device=self._device) + + # DOF state + self._dof_positions = torch.randn(self._count, self._num_dofs, device=self._device) + self._dof_velocities = torch.randn(self._count, self._num_dofs, device=self._device) + self._dof_projected_joint_forces = torch.randn(self._count, self._num_dofs, device=self._device) + + # DOF properties - stored on CPU (PhysX requirement) + self._dof_limits = torch.randn(self._count, self._num_dofs, 2, device="cpu") + self._dof_stiffnesses = torch.rand(self._count, self._num_dofs, device="cpu") * 100 + self._dof_dampings = torch.rand(self._count, self._num_dofs, device="cpu") * 10 + self._dof_max_forces = torch.rand(self._count, self._num_dofs, device="cpu") * 100 + self._dof_max_velocities = torch.rand(self._count, self._num_dofs, device="cpu") * 10 + self._dof_armatures = torch.rand(self._count, self._num_dofs, device="cpu") * 0.1 + self._dof_friction_coefficients = torch.rand(self._count, self._num_dofs, device="cpu") + self._dof_friction_properties = torch.rand(self._count, self._num_dofs, 3, device="cpu") + + # Mass properties - stored on CPU (PhysX requirement) + self._masses = torch.rand(self._count, self._num_links, device="cpu") * 10 + self._coms = torch.randn(self._count, self._num_links, 7, device="cpu") + self._coms[:, :, 3:7] = torch.nn.functional.normalize(self._coms[:, :, 3:7], dim=-1) + # Inertias: (N, L, 9) flattened format + self._inertias = torch.zeros(self._count, self._num_links, 9, device="cpu") + self._inertias[:, :, 0] = torch.rand(self._count, self._num_links) # [0,0] + self._inertias[:, :, 4] = torch.rand(self._count, self._num_links) # [1,1] + self._inertias[:, :, 8] = torch.rand(self._count, self._num_links) # [2,2] diff --git a/source/isaaclab_physx/isaaclab_physx/test/mock_interfaces/views/mock_rigid_body_view.py b/source/isaaclab_physx/isaaclab_physx/test/mock_interfaces/views/mock_rigid_body_view.py new file mode 100644 index 00000000000..a7698cd7af5 --- /dev/null +++ b/source/isaaclab_physx/isaaclab_physx/test/mock_interfaces/views/mock_rigid_body_view.py @@ -0,0 +1,360 @@ +# 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 + +"""Mock implementation of PhysX RigidBodyView.""" + +from __future__ import annotations + +import torch + + +class MockRigidBodyView: + """Mock implementation of physx.RigidBodyView for unit testing. + + This class mimics the interface of the PhysX TensorAPI RigidBodyView, + allowing tests to run without Isaac Sim or GPU simulation. + + Data Shapes: + - transforms: (N, 7) - [pos(3), quat_xyzw(4)] + - velocities: (N, 6) - [lin_vel(3), ang_vel(3)] + - accelerations: (N, 6) - [lin_acc(3), ang_acc(3)] + - masses: (N, 1) + - coms: (N, 7) - center of mass [pos(3), quat_xyzw(4)] + - inertias: (N, 9) - flattened 3x3 inertia matrix (row-major) + """ + + def __init__( + self, + count: int = 1, + prim_paths: list[str] | None = None, + device: str = "cpu", + ): + """Initialize the mock rigid body view. + + Args: + count: Number of rigid body instances. + prim_paths: USD prim paths for each instance. + device: Device for tensor allocation ("cpu" or "cuda"). + """ + self._count = count + self._prim_paths = prim_paths or [f"/World/RigidBody_{i}" for i in range(count)] + self._device = device + self._backend = "torch" + + # Internal state (lazily initialized) + self._transforms: torch.Tensor | None = None + self._velocities: torch.Tensor | None = None + self._accelerations: torch.Tensor | None = None + self._masses: torch.Tensor | None = None + self._coms: torch.Tensor | None = None + self._inertias: torch.Tensor | None = None + + # -- Helper Methods -- + + def _check_cpu_tensor(self, tensor: torch.Tensor, name: str) -> None: + """Check that tensor is on CPU, raise RuntimeError if on GPU. + + This mimics PhysX behavior where body properties must be on CPU. + """ + if tensor.is_cuda: + raise RuntimeError( + f"Expected CPU tensor for {name}, but got tensor on {tensor.device}. " + "Body properties must be set with CPU tensors." + ) + + # -- Properties -- + + @property + def count(self) -> int: + """Number of rigid body instances.""" + return self._count + + @property + def prim_paths(self) -> list[str]: + """USD prim paths for each instance.""" + return self._prim_paths + + # -- Getters -- + + def get_transforms(self) -> torch.Tensor: + """Get world transforms of all rigid bodies. + + Returns: + Tensor of shape (N, 7) with [pos_x, pos_y, pos_z, quat_x, quat_y, quat_z, quat_w]. + """ + if self._transforms is None: + # Default: origin position with identity quaternion (xyzw format) + self._transforms = torch.zeros(self._count, 7, device=self._device) + self._transforms[:, 6] = 1.0 # w=1 for identity quaternion + return self._transforms.clone() + + def get_velocities(self) -> torch.Tensor: + """Get velocities of all rigid bodies. + + Returns: + Tensor of shape (N, 6) with [lin_vel(3), ang_vel(3)]. + """ + if self._velocities is None: + self._velocities = torch.zeros(self._count, 6, device=self._device) + return self._velocities.clone() + + def get_accelerations(self) -> torch.Tensor: + """Get accelerations of all rigid bodies. + + Returns: + Tensor of shape (N, 6) with [lin_acc(3), ang_acc(3)]. + """ + if self._accelerations is None: + self._accelerations = torch.zeros(self._count, 6, device=self._device) + return self._accelerations.clone() + + def get_masses(self) -> torch.Tensor: + """Get masses of all rigid bodies. + + Returns: + Tensor of shape (N, 1) with mass values. Always on CPU. + """ + if self._masses is None: + self._masses = torch.ones(self._count, 1, device="cpu") + return self._masses.clone() + + def get_coms(self) -> torch.Tensor: + """Get centers of mass of all rigid bodies. + + Returns: + Tensor of shape (N, 7) with [pos(3), quat_xyzw(4)]. Always on CPU. + """ + if self._coms is None: + # Default: local origin with identity quaternion - stored on CPU + self._coms = torch.zeros(self._count, 7, device="cpu") + self._coms[:, 6] = 1.0 # w=1 for identity quaternion + return self._coms.clone() + + def get_inertias(self) -> torch.Tensor: + """Get inertia tensors of all rigid bodies. + + Returns: + Tensor of shape (N, 9) with flattened 3x3 inertia matrices (row-major). Always on CPU. + """ + if self._inertias is None: + # Default: identity inertia (unit sphere) - flattened [1,0,0,0,1,0,0,0,1] - stored on CPU + self._inertias = torch.zeros(self._count, 9, device="cpu") + self._inertias[:, 0] = 1.0 # [0,0] + self._inertias[:, 4] = 1.0 # [1,1] + self._inertias[:, 8] = 1.0 # [2,2] + return self._inertias.clone() + + # -- Setters (simulation interface) -- + + def set_transforms( + self, + transforms: torch.Tensor, + indices: torch.Tensor | None = None, + ) -> None: + """Set world transforms of rigid bodies. + + Args: + transforms: Tensor of shape (N, 7) or (len(indices), 7). + indices: Optional indices of bodies to update. + """ + transforms = transforms.to(self._device) + if self._transforms is None: + self._transforms = torch.zeros(self._count, 7, device=self._device) + self._transforms[:, 6] = 1.0 + if indices is not None: + self._transforms[indices] = transforms + else: + self._transforms = transforms + + def set_velocities( + self, + velocities: torch.Tensor, + indices: torch.Tensor | None = None, + ) -> None: + """Set velocities of rigid bodies. + + Args: + velocities: Tensor of shape (N, 6) or (len(indices), 6). + indices: Optional indices of bodies to update. + """ + velocities = velocities.to(self._device) + if self._velocities is None: + self._velocities = torch.zeros(self._count, 6, device=self._device) + if indices is not None: + self._velocities[indices] = velocities + else: + self._velocities = velocities + + def set_masses( + self, + masses: torch.Tensor, + indices: torch.Tensor | None = None, + ) -> None: + """Set masses of rigid bodies. + + Args: + masses: Tensor of shape (N, 1) or (len(indices), 1). Must be on CPU. + indices: Optional indices of bodies to update. + + Raises: + RuntimeError: If masses tensor is on GPU. + """ + self._check_cpu_tensor(masses, "masses") + if self._masses is None: + self._masses = torch.ones(self._count, 1, device="cpu") + if indices is not None: + self._masses[indices] = masses + else: + self._masses = masses + + def set_coms( + self, + coms: torch.Tensor, + indices: torch.Tensor | None = None, + ) -> None: + """Set centers of mass of rigid bodies. + + Args: + coms: Tensor of shape (N, 7) or (len(indices), 7). Must be on CPU. + indices: Optional indices of bodies to update. + + Raises: + RuntimeError: If coms tensor is on GPU. + """ + self._check_cpu_tensor(coms, "coms") + if self._coms is None: + self._coms = torch.zeros(self._count, 7, device="cpu") + self._coms[:, 6] = 1.0 + if indices is not None: + self._coms[indices] = coms + else: + self._coms = coms + + def set_inertias( + self, + inertias: torch.Tensor, + indices: torch.Tensor | None = None, + ) -> None: + """Set inertia tensors of rigid bodies. + + Args: + inertias: Tensor of shape (N, 9) or (len(indices), 9) - flattened 3x3 matrices. Must be on CPU. + indices: Optional indices of bodies to update. + + Raises: + RuntimeError: If inertias tensor is on GPU. + """ + self._check_cpu_tensor(inertias, "inertias") + if self._inertias is None: + self._inertias = torch.zeros(self._count, 9, device="cpu") + self._inertias[:, 0] = 1.0 + self._inertias[:, 4] = 1.0 + self._inertias[:, 8] = 1.0 + if indices is not None: + self._inertias[indices] = inertias + else: + self._inertias = inertias + + # -- Mock setters (direct test data injection) -- + + def set_mock_transforms(self, transforms: torch.Tensor) -> None: + """Set mock transform data directly for testing. + + Args: + transforms: Tensor of shape (N, 7). + """ + self._transforms = transforms.to(self._device) + + def set_mock_velocities(self, velocities: torch.Tensor) -> None: + """Set mock velocity data directly for testing. + + Args: + velocities: Tensor of shape (N, 6). + """ + self._velocities = velocities.to(self._device) + + def set_mock_accelerations(self, accelerations: torch.Tensor) -> None: + """Set mock acceleration data directly for testing. + + Args: + accelerations: Tensor of shape (N, 6). + """ + self._accelerations = accelerations.to(self._device) + + def set_mock_masses(self, masses: torch.Tensor) -> None: + """Set mock mass data directly for testing. + + Args: + masses: Tensor of shape (N, 1). + """ + self._masses = masses.to(self._device) + + def set_mock_coms(self, coms: torch.Tensor) -> None: + """Set mock center of mass data directly for testing. + + Args: + coms: Tensor of shape (N, 7). + """ + self._coms = coms.to(self._device) + + def set_mock_inertias(self, inertias: torch.Tensor) -> None: + """Set mock inertia data directly for testing. + + Args: + inertias: Tensor of shape (N, 9) - flattened 3x3 matrices. + """ + self._inertias = inertias.to(self._device) + + # -- Actions (no-op for testing) -- + + def apply_forces_and_torques_at_position( + self, + forces: torch.Tensor | None = None, + torques: torch.Tensor | None = None, + positions: torch.Tensor | None = None, + indices: torch.Tensor | None = None, + is_global: bool = True, + ) -> None: + """Apply forces and torques at positions (no-op in mock). + + Args: + forces: Forces to apply, shape (N, 3) or (len(indices), 3). + torques: Torques to apply, shape (N, 3) or (len(indices), 3). + positions: Positions to apply forces at, shape (N, 3) or (len(indices), 3). + indices: Optional indices of bodies to apply to. + is_global: Whether forces/torques are in global frame. + """ + pass # No-op for mock + + # -- Convenience method for benchmarking -- + + def set_random_mock_data(self) -> None: + """Set all internal state to random values for benchmarking. + + This method initializes all mock data with random values, + useful for benchmarking where the actual values don't matter. + """ + # Transforms with normalized quaternions - on device + self._transforms = torch.randn(self._count, 7, device=self._device) + self._transforms[:, 3:7] = torch.nn.functional.normalize(self._transforms[:, 3:7], dim=-1) + + # Velocities and accelerations - on device + self._velocities = torch.randn(self._count, 6, device=self._device) + self._accelerations = torch.randn(self._count, 6, device=self._device) + + # Mass properties - stored on CPU (PhysX requirement) + self._masses = torch.rand(self._count, 1, device="cpu") * 10 + + # Center of mass with normalized quaternions - stored on CPU (PhysX requirement) + self._coms = torch.randn(self._count, 7, device="cpu") + self._coms[:, 3:7] = torch.nn.functional.normalize(self._coms[:, 3:7], dim=-1) + + # Inertia tensors (positive definite diagonal) - flattened (N, 9) - stored on CPU (PhysX requirement) + # Create diagonal inertia matrices and flatten + diag_values = torch.rand(self._count, 3, device="cpu") + 0.1 # Ensure positive + self._inertias = torch.zeros(self._count, 9, device="cpu") + self._inertias[:, 0] = diag_values[:, 0] # [0,0] + self._inertias[:, 4] = diag_values[:, 1] # [1,1] + self._inertias[:, 8] = diag_values[:, 2] # [2,2] diff --git a/source/isaaclab_physx/isaaclab_physx/test/mock_interfaces/views/mock_rigid_contact_view.py b/source/isaaclab_physx/isaaclab_physx/test/mock_interfaces/views/mock_rigid_contact_view.py new file mode 100644 index 00000000000..9bb2821a6ba --- /dev/null +++ b/source/isaaclab_physx/isaaclab_physx/test/mock_interfaces/views/mock_rigid_contact_view.py @@ -0,0 +1,264 @@ +# 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 + +"""Mock implementation of PhysX RigidContactView.""" + +from __future__ import annotations + +import torch + + +class MockRigidContactView: + """Mock implementation of physx.RigidContactView for unit testing. + + This class mimics the interface of the PhysX TensorAPI RigidContactView, + allowing tests to run without Isaac Sim or GPU simulation. + + Data Shapes: + - net_contact_forces: (N*B, 3) - flattened net forces + - contact_force_matrix: (N*B, F, 3) - per-filter forces + - contact_data: tuple of 6 tensors (positions, normals, impulses, separations, num_found, patch_id) + - friction_data: tuple of 4 tensors (friction_forces, friction_impulses, friction_points, patch_id) + + Where: + - N = count (number of instances) + - B = num_bodies (bodies per instance) + - F = filter_count (number of filter bodies) + """ + + def __init__( + self, + count: int = 1, + num_bodies: int = 1, + filter_count: int = 0, + max_contact_data_count: int = 16, + device: str = "cpu", + ): + """Initialize the mock rigid contact view. + + Args: + count: Number of instances. + num_bodies: Number of bodies per instance. + filter_count: Number of filter bodies for contact filtering. + max_contact_data_count: Maximum number of contact data points. + device: Device for tensor allocation ("cpu" or "cuda"). + """ + self._count = count + self._num_bodies = num_bodies + self._filter_count = filter_count + self._max_contact_data_count = max_contact_data_count + self._device = device + + # Total number of bodies (flattened) + self._total_bodies = count * num_bodies + + # Internal state (lazily initialized) + self._net_contact_forces: torch.Tensor | None = None + self._contact_force_matrix: torch.Tensor | None = None + self._contact_positions: torch.Tensor | None = None + self._contact_normals: torch.Tensor | None = None + self._contact_impulses: torch.Tensor | None = None + self._contact_separations: torch.Tensor | None = None + self._contact_num_found: torch.Tensor | None = None + self._contact_patch_id: torch.Tensor | None = None + self._friction_forces: torch.Tensor | None = None + self._friction_impulses: torch.Tensor | None = None + self._friction_points: torch.Tensor | None = None + self._friction_patch_id: torch.Tensor | None = None + + # -- Properties -- + + @property + def filter_count(self) -> int: + """Number of filter bodies for contact filtering.""" + return self._filter_count + + @property + def count(self) -> int: + """Number of instances.""" + return self._count + + @property + def num_bodies(self) -> int: + """Number of bodies per instance.""" + return self._num_bodies + + # -- Getters -- + + def get_net_contact_forces(self, dt: float) -> torch.Tensor: + """Get net contact forces on all bodies. + + Args: + dt: Physics timestep (unused in mock, but required for API compatibility). + + Returns: + Tensor of shape (N*B, 3) with net forces per body. + """ + if self._net_contact_forces is None: + self._net_contact_forces = torch.zeros(self._total_bodies, 3, device=self._device) + return self._net_contact_forces.clone() + + def get_contact_force_matrix(self, dt: float) -> torch.Tensor: + """Get contact force matrix (per-filter forces). + + Args: + dt: Physics timestep (unused in mock, but required for API compatibility). + + Returns: + Tensor of shape (N*B, F, 3) with forces per body per filter. + """ + if self._contact_force_matrix is None: + self._contact_force_matrix = torch.zeros(self._total_bodies, self._filter_count, 3, device=self._device) + return self._contact_force_matrix.clone() + + def get_contact_data( + self, dt: float + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Get detailed contact data. + + Args: + dt: Physics timestep (unused in mock, but required for API compatibility). + + Returns: + Tuple of 6 tensors: + - contact_positions: (N*B, max_contacts, 3) - contact point positions + - contact_normals: (N*B, max_contacts, 3) - contact normals + - contact_impulses: (N*B, max_contacts, 3) - contact impulses + - contact_separations: (N*B, max_contacts) - contact separations + - contact_num_found: (N*B,) - number of contacts found + - contact_patch_id: (N*B, max_contacts) - contact patch IDs + """ + max_contacts = self._max_contact_data_count + + if self._contact_positions is None: + self._contact_positions = torch.zeros(self._total_bodies, max_contacts, 3, device=self._device) + if self._contact_normals is None: + self._contact_normals = torch.zeros(self._total_bodies, max_contacts, 3, device=self._device) + if self._contact_impulses is None: + self._contact_impulses = torch.zeros(self._total_bodies, max_contacts, 3, device=self._device) + if self._contact_separations is None: + self._contact_separations = torch.zeros(self._total_bodies, max_contacts, device=self._device) + if self._contact_num_found is None: + self._contact_num_found = torch.zeros(self._total_bodies, dtype=torch.int32, device=self._device) + if self._contact_patch_id is None: + self._contact_patch_id = torch.zeros( + self._total_bodies, max_contacts, dtype=torch.int32, device=self._device + ) + + return ( + self._contact_positions.clone(), + self._contact_normals.clone(), + self._contact_impulses.clone(), + self._contact_separations.clone(), + self._contact_num_found.clone(), + self._contact_patch_id.clone(), + ) + + def get_friction_data(self, dt: float) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Get friction data. + + Args: + dt: Physics timestep (unused in mock, but required for API compatibility). + + Returns: + Tuple of 4 tensors: + - friction_forces: (N*B, max_contacts, 3) - friction forces + - friction_impulses: (N*B, max_contacts, 3) - friction impulses + - friction_points: (N*B, max_contacts, 3) - friction application points + - friction_patch_id: (N*B, max_contacts) - friction patch IDs + """ + max_contacts = self._max_contact_data_count + + if self._friction_forces is None: + self._friction_forces = torch.zeros(self._total_bodies, max_contacts, 3, device=self._device) + if self._friction_impulses is None: + self._friction_impulses = torch.zeros(self._total_bodies, max_contacts, 3, device=self._device) + if self._friction_points is None: + self._friction_points = torch.zeros(self._total_bodies, max_contacts, 3, device=self._device) + if self._friction_patch_id is None: + self._friction_patch_id = torch.zeros( + self._total_bodies, max_contacts, dtype=torch.int32, device=self._device + ) + + return ( + self._friction_forces.clone(), + self._friction_impulses.clone(), + self._friction_points.clone(), + self._friction_patch_id.clone(), + ) + + # -- Mock setters (direct test data injection) -- + + def set_mock_net_contact_forces(self, forces: torch.Tensor) -> None: + """Set mock net contact force data directly for testing. + + Args: + forces: Tensor of shape (N*B, 3). + """ + self._net_contact_forces = forces.to(self._device) + + def set_mock_contact_force_matrix(self, matrix: torch.Tensor) -> None: + """Set mock contact force matrix data directly for testing. + + Args: + matrix: Tensor of shape (N*B, F, 3). + """ + self._contact_force_matrix = matrix.to(self._device) + + def set_mock_contact_data( + self, + positions: torch.Tensor | None = None, + normals: torch.Tensor | None = None, + impulses: torch.Tensor | None = None, + separations: torch.Tensor | None = None, + num_found: torch.Tensor | None = None, + patch_id: torch.Tensor | None = None, + ) -> None: + """Set mock contact data directly for testing. + + Args: + positions: Contact positions, shape (N*B, max_contacts, 3). + normals: Contact normals, shape (N*B, max_contacts, 3). + impulses: Contact impulses, shape (N*B, max_contacts, 3). + separations: Contact separations, shape (N*B, max_contacts). + num_found: Number of contacts found, shape (N*B,). + patch_id: Contact patch IDs, shape (N*B, max_contacts). + """ + if positions is not None: + self._contact_positions = positions.to(self._device) + if normals is not None: + self._contact_normals = normals.to(self._device) + if impulses is not None: + self._contact_impulses = impulses.to(self._device) + if separations is not None: + self._contact_separations = separations.to(self._device) + if num_found is not None: + self._contact_num_found = num_found.to(self._device) + if patch_id is not None: + self._contact_patch_id = patch_id.to(self._device) + + def set_mock_friction_data( + self, + forces: torch.Tensor | None = None, + impulses: torch.Tensor | None = None, + points: torch.Tensor | None = None, + patch_id: torch.Tensor | None = None, + ) -> None: + """Set mock friction data directly for testing. + + Args: + forces: Friction forces, shape (N*B, max_contacts, 3). + impulses: Friction impulses, shape (N*B, max_contacts, 3). + points: Friction application points, shape (N*B, max_contacts, 3). + patch_id: Friction patch IDs, shape (N*B, max_contacts). + """ + if forces is not None: + self._friction_forces = forces.to(self._device) + if impulses is not None: + self._friction_impulses = impulses.to(self._device) + if points is not None: + self._friction_points = points.to(self._device) + if patch_id is not None: + self._friction_patch_id = patch_id.to(self._device) diff --git a/source/isaaclab_physx/test/test_mock_interfaces/__init__.py b/source/isaaclab_physx/test/test_mock_interfaces/__init__.py new file mode 100644 index 00000000000..c360d6bd70c --- /dev/null +++ b/source/isaaclab_physx/test/test_mock_interfaces/__init__.py @@ -0,0 +1,6 @@ +# 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 + +"""Tests for mock PhysX interfaces.""" diff --git a/source/isaaclab_physx/test/test_mock_interfaces/test_mock_articulation_view.py b/source/isaaclab_physx/test/test_mock_interfaces/test_mock_articulation_view.py new file mode 100644 index 00000000000..76f25f9c23a --- /dev/null +++ b/source/isaaclab_physx/test/test_mock_interfaces/test_mock_articulation_view.py @@ -0,0 +1,332 @@ +# 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 + +"""Tests for MockArticulationView.""" + +import pytest +import torch +from isaaclab_physx.test.mock_interfaces.factories import ( + create_mock_articulation_view, + create_mock_humanoid_view, + create_mock_quadruped_view, +) +from isaaclab_physx.test.mock_interfaces.views import MockArticulationView + + +class TestMockArticulationViewInit: + """Tests for MockArticulationView initialization.""" + + def test_default_init(self): + """Test default initialization.""" + view = MockArticulationView() + assert view.count == 1 + assert view.shared_metatype.dof_count == 1 + assert view.shared_metatype.link_count == 2 + + def test_custom_configuration(self): + """Test initialization with custom configuration.""" + view = MockArticulationView( + count=4, + num_dofs=12, + num_links=13, + fixed_base=True, + ) + assert view.count == 4 + assert view.shared_metatype.dof_count == 12 + assert view.shared_metatype.link_count == 13 + assert view.shared_metatype.fixed_base is True + + def test_custom_names(self): + """Test initialization with custom DOF and link names.""" + dof_names = ["joint_0", "joint_1"] + link_names = ["base", "link_1", "link_2"] + view = MockArticulationView( + num_dofs=2, + num_links=3, + dof_names=dof_names, + link_names=link_names, + ) + assert view.shared_metatype.dof_names == dof_names + assert view.shared_metatype.link_names == link_names + + def test_tendon_properties(self): + """Test tendon properties are zero.""" + view = MockArticulationView() + assert view.max_fixed_tendons == 0 + assert view.max_spatial_tendons == 0 + + +class TestMockArticulationViewRootGetters: + """Tests for MockArticulationView root getter methods.""" + + @pytest.fixture + def view(self): + """Create a view with 4 instances, 12 DOFs, 13 links.""" + return MockArticulationView(count=4, num_dofs=12, num_links=13, device="cpu") + + def test_get_root_transforms_shape(self, view): + """Test root transforms shape.""" + transforms = view.get_root_transforms() + assert transforms.shape == (4, 7) + + def test_get_root_transforms_default_quaternion(self, view): + """Test that default quaternion is identity (xyzw format).""" + transforms = view.get_root_transforms() + assert torch.allclose(transforms[:, 6], torch.ones(4)) # w = 1 + + def test_get_root_velocities_shape(self, view): + """Test root velocities shape.""" + velocities = view.get_root_velocities() + assert velocities.shape == (4, 6) + + +class TestMockArticulationViewLinkGetters: + """Tests for MockArticulationView link getter methods.""" + + @pytest.fixture + def view(self): + """Create a view with 4 instances, 12 DOFs, 13 links.""" + return MockArticulationView(count=4, num_dofs=12, num_links=13, device="cpu") + + def test_get_link_transforms_shape(self, view): + """Test link transforms shape.""" + transforms = view.get_link_transforms() + assert transforms.shape == (4, 13, 7) + + def test_get_link_velocities_shape(self, view): + """Test link velocities shape.""" + velocities = view.get_link_velocities() + assert velocities.shape == (4, 13, 6) + + +class TestMockArticulationViewDOFGetters: + """Tests for MockArticulationView DOF getter methods.""" + + @pytest.fixture + def view(self): + """Create a view with 4 instances, 12 DOFs, 13 links.""" + return MockArticulationView(count=4, num_dofs=12, num_links=13, device="cpu") + + def test_get_dof_positions_shape(self, view): + """Test DOF positions shape.""" + positions = view.get_dof_positions() + assert positions.shape == (4, 12) + + def test_get_dof_velocities_shape(self, view): + """Test DOF velocities shape.""" + velocities = view.get_dof_velocities() + assert velocities.shape == (4, 12) + + def test_get_dof_projected_joint_forces_shape(self, view): + """Test projected joint forces shape.""" + forces = view.get_dof_projected_joint_forces() + assert forces.shape == (4, 12) + + def test_get_dof_limits_shape(self, view): + """Test DOF limits shape.""" + limits = view.get_dof_limits() + assert limits.shape == (4, 12, 2) + + def test_get_dof_limits_default_values(self, view): + """Test that default limits are infinite.""" + limits = view.get_dof_limits() + assert torch.all(limits[:, :, 0] == float("-inf")) # lower + assert torch.all(limits[:, :, 1] == float("inf")) # upper + + def test_get_dof_stiffnesses_shape(self, view): + """Test DOF stiffnesses shape.""" + stiffnesses = view.get_dof_stiffnesses() + assert stiffnesses.shape == (4, 12) + + def test_get_dof_dampings_shape(self, view): + """Test DOF dampings shape.""" + dampings = view.get_dof_dampings() + assert dampings.shape == (4, 12) + + def test_get_dof_max_forces_shape(self, view): + """Test DOF max forces shape.""" + max_forces = view.get_dof_max_forces() + assert max_forces.shape == (4, 12) + + def test_get_dof_max_velocities_shape(self, view): + """Test DOF max velocities shape.""" + max_velocities = view.get_dof_max_velocities() + assert max_velocities.shape == (4, 12) + + def test_get_dof_armatures_shape(self, view): + """Test DOF armatures shape.""" + armatures = view.get_dof_armatures() + assert armatures.shape == (4, 12) + + def test_get_dof_friction_coefficients_shape(self, view): + """Test DOF friction coefficients shape.""" + friction = view.get_dof_friction_coefficients() + assert friction.shape == (4, 12) + + +class TestMockArticulationViewMassGetters: + """Tests for MockArticulationView mass property getter methods.""" + + @pytest.fixture + def view(self): + """Create a view with 4 instances, 12 DOFs, 13 links.""" + return MockArticulationView(count=4, num_dofs=12, num_links=13, device="cpu") + + def test_get_masses_shape(self, view): + """Test masses shape.""" + masses = view.get_masses() + assert masses.shape == (4, 13) + + def test_get_coms_shape(self, view): + """Test centers of mass shape.""" + coms = view.get_coms() + assert coms.shape == (4, 13, 7) + + def test_get_inertias_shape(self, view): + """Test inertias shape.""" + inertias = view.get_inertias() + assert inertias.shape == (4, 13, 9) + + +class TestMockArticulationViewSetters: + """Tests for MockArticulationView setter methods.""" + + @pytest.fixture + def view(self): + """Create a view with 4 instances, 12 DOFs, 13 links.""" + return MockArticulationView(count=4, num_dofs=12, num_links=13, device="cpu") + + def test_set_root_transforms(self, view): + """Test setting root transforms.""" + new_data = torch.randn(4, 7) + view.set_root_transforms(new_data) + result = view.get_root_transforms() + assert torch.allclose(result, new_data) + + def test_set_dof_positions(self, view): + """Test setting DOF positions.""" + new_data = torch.randn(4, 12) + view.set_dof_positions(new_data) + result = view.get_dof_positions() + assert torch.allclose(result, new_data) + + def test_set_dof_positions_with_indices(self, view): + """Test setting DOF positions with indices.""" + new_data = torch.randn(2, 12) + indices = torch.tensor([0, 2]) + view.set_dof_positions(new_data, indices=indices) + result = view.get_dof_positions() + assert torch.allclose(result[0], new_data[0]) + assert torch.allclose(result[2], new_data[1]) + + def test_set_dof_velocities(self, view): + """Test setting DOF velocities.""" + new_data = torch.randn(4, 12) + view.set_dof_velocities(new_data) + result = view.get_dof_velocities() + assert torch.allclose(result, new_data) + + def test_set_dof_limits(self, view): + """Test setting DOF limits.""" + new_data = torch.randn(4, 12, 2) + view.set_dof_limits(new_data) + result = view.get_dof_limits() + assert torch.allclose(result, new_data) + + def test_set_dof_stiffnesses(self, view): + """Test setting DOF stiffnesses.""" + new_data = torch.randn(4, 12).abs() + view.set_dof_stiffnesses(new_data) + result = view.get_dof_stiffnesses() + assert torch.allclose(result, new_data) + + +class TestMockArticulationViewNoopSetters: + """Tests for MockArticulationView no-op setter methods.""" + + @pytest.fixture + def view(self): + """Create a view with 4 instances, 12 DOFs.""" + return MockArticulationView(count=4, num_dofs=12, device="cpu") + + def test_set_dof_position_targets_noop(self, view): + """Test that set_dof_position_targets is a no-op.""" + view.set_dof_position_targets(torch.randn(4, 12)) + + def test_set_dof_velocity_targets_noop(self, view): + """Test that set_dof_velocity_targets is a no-op.""" + view.set_dof_velocity_targets(torch.randn(4, 12)) + + def test_set_dof_actuation_forces_noop(self, view): + """Test that set_dof_actuation_forces is a no-op.""" + view.set_dof_actuation_forces(torch.randn(4, 12)) + + +class TestMockArticulationViewMockSetters: + """Tests for MockArticulationView mock setter methods.""" + + @pytest.fixture + def view(self): + """Create a view with 4 instances, 12 DOFs, 13 links.""" + return MockArticulationView(count=4, num_dofs=12, num_links=13, device="cpu") + + def test_set_mock_root_transforms(self, view): + """Test mock root transform setter.""" + mock_data = torch.randn(4, 7) + view.set_mock_root_transforms(mock_data) + result = view.get_root_transforms() + assert torch.allclose(result, mock_data) + + def test_set_mock_link_transforms(self, view): + """Test mock link transform setter.""" + mock_data = torch.randn(4, 13, 7) + view.set_mock_link_transforms(mock_data) + result = view.get_link_transforms() + assert torch.allclose(result, mock_data) + + def test_set_mock_dof_positions(self, view): + """Test mock DOF position setter.""" + mock_data = torch.randn(4, 12) + view.set_mock_dof_positions(mock_data) + result = view.get_dof_positions() + assert torch.allclose(result, mock_data) + + def test_set_mock_dof_velocities(self, view): + """Test mock DOF velocity setter.""" + mock_data = torch.randn(4, 12) + view.set_mock_dof_velocities(mock_data) + result = view.get_dof_velocities() + assert torch.allclose(result, mock_data) + + +class TestMockArticulationViewFactories: + """Tests for articulation view factory functions.""" + + def test_create_mock_articulation_view_basic(self): + """Test basic factory usage.""" + view = create_mock_articulation_view(count=4, num_dofs=12, num_links=13) + assert view.count == 4 + assert view.shared_metatype.dof_count == 12 + assert view.shared_metatype.link_count == 13 + + def test_create_mock_quadruped_view(self): + """Test quadruped factory.""" + view = create_mock_quadruped_view(count=4) + assert view.count == 4 + assert view.shared_metatype.dof_count == 12 + assert view.shared_metatype.link_count == 13 + assert view.shared_metatype.fixed_base is False + assert "FL_hip_joint" in view.shared_metatype.dof_names + assert "base" in view.shared_metatype.link_names + + def test_create_mock_humanoid_view(self): + """Test humanoid factory.""" + view = create_mock_humanoid_view(count=2) + assert view.count == 2 + assert view.shared_metatype.dof_count == 21 + assert view.shared_metatype.link_count == 22 + assert view.shared_metatype.fixed_base is False + assert "left_elbow" in view.shared_metatype.dof_names + assert "pelvis" in view.shared_metatype.link_names diff --git a/source/isaaclab_physx/test/test_mock_interfaces/test_mock_rigid_body_view.py b/source/isaaclab_physx/test/test_mock_interfaces/test_mock_rigid_body_view.py new file mode 100644 index 00000000000..014423ee481 --- /dev/null +++ b/source/isaaclab_physx/test/test_mock_interfaces/test_mock_rigid_body_view.py @@ -0,0 +1,217 @@ +# 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 + +"""Tests for MockRigidBodyView.""" + +import pytest +import torch +from isaaclab_physx.test.mock_interfaces.factories import create_mock_rigid_body_view +from isaaclab_physx.test.mock_interfaces.views import MockRigidBodyView + + +class TestMockRigidBodyViewInit: + """Tests for MockRigidBodyView initialization.""" + + def test_default_init(self): + """Test default initialization.""" + view = MockRigidBodyView() + assert view.count == 1 + assert len(view.prim_paths) == 1 + assert view._backend == "torch" + + def test_custom_count(self): + """Test initialization with custom count.""" + view = MockRigidBodyView(count=4) + assert view.count == 4 + assert len(view.prim_paths) == 4 + + def test_custom_prim_paths(self): + """Test initialization with custom prim paths.""" + paths = ["/World/Body_A", "/World/Body_B"] + view = MockRigidBodyView(count=2, prim_paths=paths) + assert view.prim_paths == paths + + +class TestMockRigidBodyViewGetters: + """Tests for MockRigidBodyView getter methods.""" + + @pytest.fixture + def view(self): + """Create a view with 4 instances.""" + return MockRigidBodyView(count=4, device="cpu") + + def test_get_transforms_shape(self, view): + """Test transforms shape.""" + transforms = view.get_transforms() + assert transforms.shape == (4, 7) + + def test_get_transforms_default_quaternion(self, view): + """Test that default quaternion is identity (xyzw format).""" + transforms = view.get_transforms() + # xyzw format: [x, y, z, w] = [0, 0, 0, 1] for identity + assert torch.allclose(transforms[:, 3:6], torch.zeros(4, 3)) # xyz = 0 + assert torch.allclose(transforms[:, 6], torch.ones(4)) # w = 1 + + def test_get_velocities_shape(self, view): + """Test velocities shape.""" + velocities = view.get_velocities() + assert velocities.shape == (4, 6) + + def test_get_accelerations_shape(self, view): + """Test accelerations shape.""" + accelerations = view.get_accelerations() + assert accelerations.shape == (4, 6) + + def test_get_masses_shape(self, view): + """Test masses shape.""" + masses = view.get_masses() + assert masses.shape == (4, 1) + + def test_get_masses_default_value(self, view): + """Test that default mass is 1.""" + masses = view.get_masses() + assert torch.allclose(masses, torch.ones(4, 1)) + + def test_get_coms_shape(self, view): + """Test centers of mass shape.""" + coms = view.get_coms() + assert coms.shape == (4, 7) + + def test_get_inertias_shape(self, view): + """Test inertias shape.""" + inertias = view.get_inertias() + assert inertias.shape == (4, 9) + + def test_get_inertias_default_value(self, view): + """Test that default inertia is identity.""" + inertias = view.get_inertias() + expected = torch.eye(3).repeat(4, 1).reshape(4, 9) + assert torch.allclose(inertias, expected) + + def test_getters_return_clones(self, view): + """Test that getters return clones, not references.""" + transforms1 = view.get_transforms() + transforms1[0, 0] = 999.0 + transforms2 = view.get_transforms() + assert transforms2[0, 0] != 999.0 + + +class TestMockRigidBodyViewSetters: + """Tests for MockRigidBodyView setter methods.""" + + @pytest.fixture + def view(self): + """Create a view with 4 instances.""" + return MockRigidBodyView(count=4, device="cpu") + + def test_set_transforms(self, view): + """Test setting transforms.""" + new_transforms = torch.randn(4, 7) + view.set_transforms(new_transforms) + result = view.get_transforms() + assert torch.allclose(result, new_transforms) + + def test_set_transforms_with_indices(self, view): + """Test setting transforms with indices.""" + new_transforms = torch.randn(2, 7) + indices = torch.tensor([0, 2]) + view.set_transforms(new_transforms, indices=indices) + result = view.get_transforms() + assert torch.allclose(result[0], new_transforms[0]) + assert torch.allclose(result[2], new_transforms[1]) + + def test_set_velocities(self, view): + """Test setting velocities.""" + new_velocities = torch.randn(4, 6) + view.set_velocities(new_velocities) + result = view.get_velocities() + assert torch.allclose(result, new_velocities) + + def test_set_masses(self, view): + """Test setting masses.""" + new_masses = torch.randn(4, 1, 1).abs() + view.set_masses(new_masses) + result = view.get_masses() + assert torch.allclose(result, new_masses) + + +class TestMockRigidBodyViewMockSetters: + """Tests for MockRigidBodyView mock setter methods.""" + + @pytest.fixture + def view(self): + """Create a view with 4 instances.""" + return MockRigidBodyView(count=4, device="cpu") + + def test_set_mock_transforms(self, view): + """Test mock transform setter.""" + mock_data = torch.randn(4, 7) + view.set_mock_transforms(mock_data) + result = view.get_transforms() + assert torch.allclose(result, mock_data) + + def test_set_mock_velocities(self, view): + """Test mock velocity setter.""" + mock_data = torch.randn(4, 6) + view.set_mock_velocities(mock_data) + result = view.get_velocities() + assert torch.allclose(result, mock_data) + + def test_set_mock_accelerations(self, view): + """Test mock acceleration setter.""" + mock_data = torch.randn(4, 6) + view.set_mock_accelerations(mock_data) + result = view.get_accelerations() + assert torch.allclose(result, mock_data) + + def test_set_mock_masses(self, view): + """Test mock mass setter.""" + mock_data = torch.randn(4, 1, 1).abs() + view.set_mock_masses(mock_data) + result = view.get_masses() + assert torch.allclose(result, mock_data) + + def test_set_mock_coms(self, view): + """Test mock center of mass setter.""" + mock_data = torch.randn(4, 1, 7) + view.set_mock_coms(mock_data) + result = view.get_coms() + assert torch.allclose(result, mock_data) + + def test_set_mock_inertias(self, view): + """Test mock inertia setter.""" + mock_data = torch.randn(4, 1, 3, 3) + view.set_mock_inertias(mock_data) + result = view.get_inertias() + assert torch.allclose(result, mock_data) + + +class TestMockRigidBodyViewActions: + """Tests for MockRigidBodyView action methods.""" + + def test_apply_forces_and_torques_at_position_noop(self): + """Test that apply_forces_and_torques_at_position is a no-op.""" + view = MockRigidBodyView(count=4) + # Should not raise + view.apply_forces_and_torques_at_position( + forces=torch.randn(4, 3), + torques=torch.randn(4, 3), + positions=torch.randn(4, 3), + ) + + +class TestMockRigidBodyViewFactory: + """Tests for create_mock_rigid_body_view factory function.""" + + def test_factory_basic(self): + """Test basic factory usage.""" + view = create_mock_rigid_body_view(count=4) + assert view.count == 4 + + def test_factory_with_prim_paths(self): + """Test factory with custom prim paths.""" + paths = ["/World/A", "/World/B"] + view = create_mock_rigid_body_view(count=2, prim_paths=paths) + assert view.prim_paths == paths diff --git a/source/isaaclab_physx/test/test_mock_interfaces/test_mock_rigid_contact_view.py b/source/isaaclab_physx/test/test_mock_interfaces/test_mock_rigid_contact_view.py new file mode 100644 index 00000000000..4c2503531da --- /dev/null +++ b/source/isaaclab_physx/test/test_mock_interfaces/test_mock_rigid_contact_view.py @@ -0,0 +1,196 @@ +# 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 + +"""Tests for MockRigidContactView.""" + +import pytest +import torch +from isaaclab_physx.test.mock_interfaces.factories import create_mock_rigid_contact_view +from isaaclab_physx.test.mock_interfaces.views import MockRigidContactView + + +class TestMockRigidContactViewInit: + """Tests for MockRigidContactView initialization.""" + + def test_default_init(self): + """Test default initialization.""" + view = MockRigidContactView() + assert view.count == 1 + assert view.num_bodies == 1 + assert view.filter_count == 0 + + def test_custom_configuration(self): + """Test initialization with custom configuration.""" + view = MockRigidContactView( + count=4, + num_bodies=5, + filter_count=3, + ) + assert view.count == 4 + assert view.num_bodies == 5 + assert view.filter_count == 3 + + +class TestMockRigidContactViewGetters: + """Tests for MockRigidContactView getter methods.""" + + @pytest.fixture + def view(self): + """Create a view with 4 instances, 5 bodies, 3 filters.""" + return MockRigidContactView(count=4, num_bodies=5, filter_count=3, device="cpu") + + def test_get_net_contact_forces_shape(self, view): + """Test net contact forces shape.""" + forces = view.get_net_contact_forces(dt=0.01) + # Shape: (N*B, 3) = (4*5, 3) = (20, 3) + assert forces.shape == (20, 3) + + def test_get_contact_force_matrix_shape(self, view): + """Test contact force matrix shape.""" + matrix = view.get_contact_force_matrix(dt=0.01) + # Shape: (N*B, F, 3) = (4*5, 3, 3) = (20, 3, 3) + assert matrix.shape == (20, 3, 3) + + def test_get_contact_data_shapes(self, view): + """Test contact data tuple shapes.""" + data = view.get_contact_data(dt=0.01) + assert len(data) == 6 + positions, normals, impulses, separations, num_found, patch_id = data + + total_bodies = 4 * 5 # count * num_bodies + max_contacts = 16 # default + + assert positions.shape == (total_bodies, max_contacts, 3) + assert normals.shape == (total_bodies, max_contacts, 3) + assert impulses.shape == (total_bodies, max_contacts, 3) + assert separations.shape == (total_bodies, max_contacts) + assert num_found.shape == (total_bodies,) + assert patch_id.shape == (total_bodies, max_contacts) + + def test_get_friction_data_shapes(self, view): + """Test friction data tuple shapes.""" + data = view.get_friction_data(dt=0.01) + assert len(data) == 4 + forces, impulses, points, patch_id = data + + total_bodies = 4 * 5 + max_contacts = 16 + + assert forces.shape == (total_bodies, max_contacts, 3) + assert impulses.shape == (total_bodies, max_contacts, 3) + assert points.shape == (total_bodies, max_contacts, 3) + assert patch_id.shape == (total_bodies, max_contacts) + + def test_getters_return_clones(self, view): + """Test that getters return clones, not references.""" + forces1 = view.get_net_contact_forces(0.01) + forces1[0, 0] = 999.0 + forces2 = view.get_net_contact_forces(0.01) + assert forces2[0, 0] != 999.0 + + +class TestMockRigidContactViewMockSetters: + """Tests for MockRigidContactView mock setter methods.""" + + @pytest.fixture + def view(self): + """Create a view with 4 instances, 5 bodies, 3 filters.""" + return MockRigidContactView(count=4, num_bodies=5, filter_count=3, device="cpu") + + def test_set_mock_net_contact_forces(self, view): + """Test mock net contact forces setter.""" + mock_data = torch.randn(20, 3) + view.set_mock_net_contact_forces(mock_data) + result = view.get_net_contact_forces(0.01) + assert torch.allclose(result, mock_data) + + def test_set_mock_contact_force_matrix(self, view): + """Test mock contact force matrix setter.""" + mock_data = torch.randn(20, 3, 3) + view.set_mock_contact_force_matrix(mock_data) + result = view.get_contact_force_matrix(0.01) + assert torch.allclose(result, mock_data) + + def test_set_mock_contact_data_partial(self, view): + """Test setting partial contact data.""" + mock_positions = torch.randn(20, 16, 3) + mock_normals = torch.randn(20, 16, 3) + view.set_mock_contact_data(positions=mock_positions, normals=mock_normals) + + positions, normals, _, _, _, _ = view.get_contact_data(0.01) + assert torch.allclose(positions, mock_positions) + assert torch.allclose(normals, mock_normals) + + def test_set_mock_contact_data_full(self, view): + """Test setting full contact data.""" + total_bodies = 20 + max_contacts = 16 + + mock_positions = torch.randn(total_bodies, max_contacts, 3) + mock_normals = torch.randn(total_bodies, max_contacts, 3) + mock_impulses = torch.randn(total_bodies, max_contacts, 3) + mock_separations = torch.randn(total_bodies, max_contacts) + mock_num_found = torch.randint(0, max_contacts, (total_bodies,), dtype=torch.int32) + mock_patch_id = torch.randint(0, 10, (total_bodies, max_contacts), dtype=torch.int32) + + view.set_mock_contact_data( + positions=mock_positions, + normals=mock_normals, + impulses=mock_impulses, + separations=mock_separations, + num_found=mock_num_found, + patch_id=mock_patch_id, + ) + + positions, normals, impulses, separations, num_found, patch_id = view.get_contact_data(0.01) + + assert torch.allclose(positions, mock_positions) + assert torch.allclose(normals, mock_normals) + assert torch.allclose(impulses, mock_impulses) + assert torch.allclose(separations, mock_separations) + assert torch.equal(num_found, mock_num_found) + assert torch.equal(patch_id, mock_patch_id) + + def test_set_mock_friction_data(self, view): + """Test setting friction data.""" + total_bodies = 20 + max_contacts = 16 + + mock_forces = torch.randn(total_bodies, max_contacts, 3) + mock_impulses = torch.randn(total_bodies, max_contacts, 3) + + view.set_mock_friction_data(forces=mock_forces, impulses=mock_impulses) + + forces, impulses, _, _ = view.get_friction_data(0.01) + assert torch.allclose(forces, mock_forces) + assert torch.allclose(impulses, mock_impulses) + + +class TestMockRigidContactViewFactory: + """Tests for create_mock_rigid_contact_view factory function.""" + + def test_factory_basic(self): + """Test basic factory usage.""" + view = create_mock_rigid_contact_view(count=4, num_bodies=5, filter_count=3) + assert view.count == 4 + assert view.num_bodies == 5 + assert view.filter_count == 3 + + def test_factory_custom_max_contacts(self): + """Test factory with custom max contact data count.""" + view = create_mock_rigid_contact_view(count=2, num_bodies=3, max_contact_data_count=32) + _, _, _, separations, _, _ = view.get_contact_data(0.01) + assert separations.shape[1] == 32 + + +class TestMockRigidContactViewZeroFilterCount: + """Tests for MockRigidContactView with zero filter count.""" + + def test_contact_force_matrix_zero_filters(self): + """Test contact force matrix with zero filters.""" + view = MockRigidContactView(count=4, num_bodies=5, filter_count=0) + matrix = view.get_contact_force_matrix(0.01) + # Shape: (N*B, F, 3) = (20, 0, 3) + assert matrix.shape == (20, 0, 3) diff --git a/source/isaaclab_physx/test/test_mock_interfaces/test_patching.py b/source/isaaclab_physx/test/test_mock_interfaces/test_patching.py new file mode 100644 index 00000000000..7c49afdef34 --- /dev/null +++ b/source/isaaclab_physx/test/test_mock_interfaces/test_patching.py @@ -0,0 +1,180 @@ +# 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 + +"""Tests for patching utilities.""" + +from isaaclab_physx.test.mock_interfaces.utils import ( + mock_articulation_view, + mock_rigid_body_view, + mock_rigid_contact_view, + patch_articulation_view, + patch_rigid_body_view, + patch_rigid_contact_view, +) + + +class TestPatchRigidBodyView: + """Tests for patch_rigid_body_view context manager.""" + + def test_basic_patching(self): + """Test basic patching functionality. + + Note: Patching works when the target module imports the class dynamically. + In tests where we import at module level, we need to import inside the + context to get the patched version. + """ + target = "isaaclab_physx.test.mock_interfaces.views.mock_rigid_body_view.MockRigidBodyView" + with patch_rigid_body_view(target, count=4): + # Import inside context to get patched version + from isaaclab_physx.test.mock_interfaces.views import mock_rigid_body_view + + view = mock_rigid_body_view.MockRigidBodyView() + assert view.count == 4 + + def test_patching_preserves_configuration(self): + """Test that patching preserves configuration.""" + target = "isaaclab_physx.test.mock_interfaces.views.mock_rigid_body_view.MockRigidBodyView" + prim_paths = ["/World/A", "/World/B", "/World/C", "/World/D"] + with patch_rigid_body_view(target, count=4, prim_paths=prim_paths): + from isaaclab_physx.test.mock_interfaces.views import mock_rigid_body_view + + view = mock_rigid_body_view.MockRigidBodyView() + assert view.prim_paths == prim_paths + + +class TestPatchArticulationView: + """Tests for patch_articulation_view context manager.""" + + def test_basic_patching(self): + """Test basic patching functionality.""" + target = "isaaclab_physx.test.mock_interfaces.views.mock_articulation_view.MockArticulationView" + with patch_articulation_view(target, count=4, num_dofs=12, num_links=13): + from isaaclab_physx.test.mock_interfaces.views import mock_articulation_view + + view = mock_articulation_view.MockArticulationView() + assert view.count == 4 + assert view.shared_metatype.dof_count == 12 + assert view.shared_metatype.link_count == 13 + + def test_patching_with_names(self): + """Test patching with custom names.""" + target = "isaaclab_physx.test.mock_interfaces.views.mock_articulation_view.MockArticulationView" + dof_names = ["joint_a", "joint_b"] + with patch_articulation_view(target, num_dofs=2, dof_names=dof_names): + from isaaclab_physx.test.mock_interfaces.views import mock_articulation_view + + view = mock_articulation_view.MockArticulationView() + assert view.shared_metatype.dof_names == dof_names + + +class TestPatchRigidContactView: + """Tests for patch_rigid_contact_view context manager.""" + + def test_basic_patching(self): + """Test basic patching functionality.""" + target = "isaaclab_physx.test.mock_interfaces.views.mock_rigid_contact_view.MockRigidContactView" + with patch_rigid_contact_view(target, count=4, num_bodies=5, filter_count=3): + from isaaclab_physx.test.mock_interfaces.views import mock_rigid_contact_view + + view = mock_rigid_contact_view.MockRigidContactView() + assert view.count == 4 + assert view.num_bodies == 5 + assert view.filter_count == 3 + + +class TestDecoratorUtilities: + """Tests for decorator utilities. + + Note: These decorators are designed for wrapping functions, not pytest tests. + Pytest inspects function signatures for fixtures, which conflicts with our + decorator pattern. We test them by manually invoking the decorated functions. + """ + + def test_mock_rigid_body_view_decorator(self): + """Test mock_rigid_body_view decorator injects mock view.""" + + @mock_rigid_body_view(count=4) + def my_function(view): + return view.count, view.get_transforms().shape + + count, shape = my_function() + assert count == 4 + assert shape == (4, 7) + + def test_mock_articulation_view_decorator(self): + """Test mock_articulation_view decorator injects mock view.""" + + @mock_articulation_view(count=4, num_dofs=12, num_links=13) + def my_function(view): + return ( + view.count, + view.shared_metatype.dof_count, + view.get_dof_positions().shape, + ) + + count, dof_count, shape = my_function() + assert count == 4 + assert dof_count == 12 + assert shape == (4, 12) + + def test_mock_rigid_contact_view_decorator(self): + """Test mock_rigid_contact_view decorator injects mock view.""" + + @mock_rigid_contact_view(count=4, num_bodies=5, filter_count=3) + def my_function(view): + return ( + view.count, + view.num_bodies, + view.get_net_contact_forces(0.01).shape, + ) + + count, num_bodies, shape = my_function() + assert count == 4 + assert num_bodies == 5 + assert shape == (20, 3) # 4 * 5 = 20 + + def test_decorator_with_prim_paths(self): + """Test decorator with custom prim paths.""" + + @mock_rigid_body_view(count=2, prim_paths=["/World/A", "/World/B"]) + def my_function(view): + return view.prim_paths + + paths = my_function() + assert paths == ["/World/A", "/World/B"] + + def test_decorator_with_fixed_base(self): + """Test decorator with fixed base.""" + + @mock_articulation_view(count=2, num_dofs=6, fixed_base=True) + def my_function(view): + return view.shared_metatype.fixed_base + + fixed_base = my_function() + assert fixed_base is True + + +class TestDecoratorFunctionPreservation: + """Tests that decorators preserve function metadata.""" + + def test_function_name_preserved(self): + """Test that function name is preserved.""" + + @mock_rigid_body_view(count=4) + def my_documented_function(view): + """This is a documented function.""" + pass + + assert my_documented_function.__name__ == "my_documented_function" + + def test_docstring_preserved(self): + """Test that docstring is preserved.""" + + @mock_rigid_body_view(count=4) + def my_documented_function(view): + """This is a documented function.""" + pass + + assert "documented function" in my_documented_function.__doc__