diff --git a/docs/source/how-to/index.rst b/docs/source/how-to/index.rst index 278e88c4b6c6..85ef6055839e 100644 --- a/docs/source/how-to/index.rst +++ b/docs/source/how-to/index.rst @@ -101,6 +101,18 @@ Isaac Lab. draw_markers +Working with Simulation Data +---------------------------- + +This guide explains how to read and write simulation state using the :class:`~isaaclab.utils.warp.ProxyArray` +dual-access wrapper that all data classes return. + +.. toctree:: + :maxdepth: 1 + + proxy_array + + Interfacing with Environments ----------------------------- diff --git a/docs/source/how-to/proxy_array.rst b/docs/source/how-to/proxy_array.rst new file mode 100644 index 000000000000..7bf0a90374f3 --- /dev/null +++ b/docs/source/how-to/proxy_array.rst @@ -0,0 +1,134 @@ +.. _how-to-torch-array: + +Working with ProxyArray +======================= + +.. currentmodule:: isaaclab.utils.warp + +Isaac Lab data classes return :class:`ProxyArray` — a lightweight, warp-first wrapper that +provides zero-copy access to simulation data as either a :class:`warp.array` or a +:class:`torch.Tensor`. + +.. note:: + + ``ProxyArray`` is inspired by the ``ProxyArray`` class from + `mujocolab/mjlab `_ (BSD-3-Clause). + The design adapts the same dual-accessor pattern to Isaac Lab's warp-based data pipeline. + + +Quick Start +~~~~~~~~~~~ + +Every property on asset and sensor data classes (e.g., ``robot.data.joint_pos``, +``sensor.data.net_forces_w``) returns a ``ProxyArray``: + +.. code-block:: python + + robot = env.scene["robot"] + + # Explicit torch access (preferred) + joint_positions = robot.data.joint_pos.torch # torch.Tensor, cached zero-copy view + gravity_proj = robot.data.projected_gravity_b.torch # torch.Tensor + + # Explicit warp access (for kernel interop) + wp_array = robot.data.joint_pos.warp # warp.array, the original buffer + + # Pass directly to warp kernels — no unwrapping needed + wp.launch(my_kernel, inputs=[robot.data.joint_pos], ...) # works via __cuda_array_interface__ + + +The ``.torch`` and ``.warp`` Accessors +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +- ``.torch`` returns a cached, zero-copy :class:`torch.Tensor` view of the underlying warp array + (via :func:`warp.to_torch`). The tensor is created on first access and reused on subsequent + calls. Since it is a zero-copy view, writes to the tensor are visible through the warp array + and vice versa. + +- ``.warp`` returns the original :class:`warp.array`. Use this when you need to pass data to + warp kernels explicitly or access warp-specific attributes (``ptr``, ``strides``, etc.). + + +Deprecation Bridge +~~~~~~~~~~~~~~~~~~ + +To ease migration, ``ProxyArray`` includes a deprecation bridge that allows existing code to +treat it as a ``torch.Tensor`` temporarily: + +.. code-block:: python + + # These still work but emit a one-time DeprecationWarning: + result = torch.sum(robot.data.joint_pos) # via __torch_function__ + value = robot.data.joint_pos[0, 3] # via __getitem__ + result = robot.data.joint_pos + 1.0 # via __add__ + legacy = wp.to_torch(robot.data.joint_pos) # temporary shim + + # Preferred (no warning): + result = torch.sum(robot.data.joint_pos.torch) + value = robot.data.joint_pos.torch[0, 3] + result = robot.data.joint_pos.torch + 1.0 + tensor = robot.data.joint_pos.torch + +The ``wp.to_torch()`` compatibility path is a temporary shim for code that was migrated +before ``ProxyArray`` exposed explicit accessors. It returns the same zero-copy tensor as +``.torch`` and emits a one-time ``DeprecationWarning``. The shim and the other deprecation +bridges will be removed in a future release. Migrate to explicit ``.torch`` access now. + + +Migrating from Isaac Lab 2.x +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +In Isaac Lab 2.x, data properties returned ``torch.Tensor`` directly. In 3.0, they return +``ProxyArray``. Append ``.torch`` to get the tensor: + +.. code-block:: python + + # BEFORE (Isaac Lab 2.x) — properties returned torch.Tensor directly + joint_pos = robot.data.joint_pos + first_contact = sensor.compute_first_contact(dt) + + # AFTER (Isaac Lab 3.0) — properties return ProxyArray + joint_pos = robot.data.joint_pos.torch + first_contact = sensor.compute_first_contact(dt).torch + +.. note:: + + Passing a ``ProxyArray`` to ``wp.to_torch()`` is temporarily supported by a compatibility + shim and returns the same cached zero-copy tensor as ``.torch``. This path emits a + one-time ``DeprecationWarning`` and will be removed in a future release. Use ``.torch`` + in new code. + + +Backend Differences +~~~~~~~~~~~~~~~~~~~ + +While the ``ProxyArray`` interface is identical across backends, the underlying data refresh +model differs: + +**PhysX (pull-to-refresh):** + Properties pull fresh data from the PhysX tensor API on first access per simulation step, + then cache the result. The underlying GPU buffers are stable and pre-allocated — the + ``ProxyArray`` wrapper is created once and reused safely across steps. + +**Newton (auto-refresh with wrapper replacement):** + The simulation automatically refreshes GPU buffers each step. On full simulation resets, + buffers may be re-created. The Newton backend creates new ``ProxyArray`` wrappers for the + new warp arrays, invalidating any previously cached torch tensors. + +In both cases, ``.torch`` always returns a view of the current simulation state for the +current step. + +.. warning:: + + Do not cache ``.torch`` results across simulation resets. After a reset (especially on + Newton), previously obtained tensors may point to stale or freed GPU memory. Always + re-access the property after a reset. + + +API Reference +~~~~~~~~~~~~~ + +.. autoclass:: ProxyArray + :members: + :undoc-members: + :no-index: diff --git a/docs/source/migration/migrating_to_isaaclab_3-0.rst b/docs/source/migration/migrating_to_isaaclab_3-0.rst index b714257797b0..f90d47496d47 100644 --- a/docs/source/migration/migrating_to_isaaclab_3-0.rst +++ b/docs/source/migration/migrating_to_isaaclab_3-0.rst @@ -889,6 +889,56 @@ Best Practices for Migration 6. **Check documentation** - Update any docs or comments that mention quaternion format. + +Using the Runtime Quaternion Access Detector +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +The quaternion finder tool above covers hard-coded values in source files, +but it cannot see quaternions that are *read* from asset/sensor data at +runtime. For those, Isaac Lab ships a +runtime detector hook on :class:`~isaaclab.utils.warp.ProxyArray` that flags +every ``.torch`` access on a ``wp.quatf``-typed property and points at the +exact call site. Use it after the source-level migration to catch the cases +the finder tool can't reach. + +Enable it by setting an environment variable before launching your script: + +.. code-block:: bash + + export WARN_ON_TORCH_QUATF_ACCESS=1 + ./isaaclab.sh -p my_script.py + +Every read of ``.torch`` on a ``ProxyArray`` whose underlying ``wp.array`` has +dtype ``wp.quatf`` then emits a :class:`UserWarning` with the message: + +.. code-block:: text + + Reading .torch on a wp.quatf-typed ProxyArray. The Isaac Lab quaternion + convention changed from (w, x, y, z) in 2.x to (x, y, z, w) in 3.x. If + your code assumes the old order, this is likely the source of incorrect + rotations. Unset WARN_ON_TORCH_QUATF_ACCESS to silence this warning. + +The warning's traceback points at the exact line that performed the access +(via ``stacklevel=2``), so you can walk through the matches in your code and +confirm each one uses the new ``(x, y, z, w)`` order. + +Typical workflow: + +1. Run a representative scene or task with the env var set. +2. Triage every warning location — check whether the call site assumes + ``(w, x, y, z)`` (Lab 2.x) or ``(x, y, z, w)`` (Lab 3.x). +3. Migrate the call sites that still expect the old order. +4. Re-run with the env var still set; the warnings should be gone (or only + come from intentionally-handled call sites). +5. Unset the env var for production runs — the detector adds an + ``os.environ`` lookup per ``.torch`` access, which is cheap but not free. + +The detector covers only ``ProxyArray.torch`` reads. Direct accesses on the +underlying ``wp.array`` (via ``ProxyArray.warp``) are not flagged, because +warp uses ``(x, y, z, w)`` natively and so a warp-side read is unaffected +by the convention change. + + API Changes ~~~~~~~~~~~ @@ -922,37 +972,36 @@ quaternions in XYZW format: - And all other quaternion utilities -Warp Backend for Asset and Sensor Data -~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +ProxyArray Backend for Asset and Sensor Data +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ -All ``.data.*`` properties on asset and sensor classes now return ``wp.array`` instead of -``torch.Tensor``. This change applies to all asset classes (:class:`~isaaclab.assets.Articulation`, +All ``.data.*`` properties on asset and sensor classes now return +:class:`~isaaclab.utils.warp.ProxyArray` instead of ``torch.Tensor``. ``ProxyArray`` wraps +the underlying ``wp.array`` and exposes explicit ``.torch`` and ``.warp`` accessors. This +change applies to all asset classes (:class:`~isaaclab.assets.Articulation`, :class:`~isaaclab.assets.RigidObject`, :class:`~isaaclab.assets.RigidObjectCollection`, :class:`~isaaclab_physx.assets.DeformableObject`) and all sensor classes (:class:`~isaaclab_physx.sensors.ContactSensor`, :class:`~isaaclab_physx.sensors.Imu`, :class:`~isaaclab_physx.sensors.Pva`, :class:`~isaaclab_physx.sensors.FrameTransformer`). -To convert back to ``torch.Tensor`` for use with PyTorch operations, wrap the property -access with ``wp.to_torch()``: +To use a data property as a ``torch.Tensor``, append ``.torch``: .. code-block:: python - import warp as wp - # Before (Isaac Lab 2.x) root_pos = robot.data.root_pos_w # torch.Tensor joint_pos = robot.data.joint_pos # torch.Tensor contact_forces = sensor.data.net_forces_w # torch.Tensor # After (Isaac Lab 3.x) - root_pos = robot.data.root_pos_w # wp.array - joint_pos = robot.data.joint_pos # wp.array - contact_forces = sensor.data.net_forces_w # wp.array + root_pos = robot.data.root_pos_w # ProxyArray + joint_pos = robot.data.joint_pos # ProxyArray + contact_forces = sensor.data.net_forces_w # ProxyArray - # To use with torch operations, wrap with wp.to_torch() - root_pos_torch = wp.to_torch(robot.data.root_pos_w) # torch.Tensor - joint_pos_torch = wp.to_torch(robot.data.joint_pos) # torch.Tensor - contact_torch = wp.to_torch(sensor.data.net_forces_w) # torch.Tensor + # To use with torch operations, access .torch + root_pos_torch = robot.data.root_pos_w.torch # torch.Tensor + joint_pos_torch = robot.data.joint_pos.torch # torch.Tensor + contact_torch = sensor.data.net_forces_w.torch # torch.Tensor Common patterns that need updating: @@ -962,19 +1011,19 @@ Common patterns that need updating: # Before: pos = robot.data.root_pos_w.clone() # After: - pos = wp.to_torch(robot.data.root_pos_w).clone() + pos = robot.data.root_pos_w.torch.clone() # Creating zero tensors with matching shape # Before: zeros = torch.zeros_like(robot.data.root_pos_w) # After: - zeros = torch.zeros_like(wp.to_torch(robot.data.root_pos_w)) + zeros = torch.zeros_like(robot.data.root_pos_w.torch) # Assertions in tests # Before: torch.testing.assert_close(robot.data.root_pos_w, expected) # After: - torch.testing.assert_close(wp.to_torch(robot.data.root_pos_w), expected) + torch.testing.assert_close(robot.data.root_pos_w.torch, expected) .. list-table:: Affected classes :header-rows: 1 @@ -1009,20 +1058,10 @@ Common patterns that need updating: .. note:: - An automated migration tool is provided at ``scripts/tools/wrap_warp_to_torch.py``. - It scans Python files for ``.data.`` accesses and wraps them with - ``wp.to_torch()``. Usage: - - .. code-block:: bash - - # Dry run (preview changes) - python scripts/tools/wrap_warp_to_torch.py path/to/your/code --dry-run - - # Apply changes in-place - python scripts/tools/wrap_warp_to_torch.py path/to/your/code - - Always review the changes after running the tool, as some accesses (e.g., those - already passed to warp-native functions) should not be wrapped. + ``wp.to_torch(proxy_array)`` is temporarily supported by a compatibility shim. It returns + the same zero-copy tensor as ``proxy_array.torch`` and emits a one-time + ``DeprecationWarning``. This shim exists for older migration code and will be removed in a + future release; prefer ``.torch`` in new code. Ray Caster Warp Backend @@ -1041,25 +1080,23 @@ RayCasterData Return Types The :attr:`~isaaclab.sensors.RayCasterData.pos_w`, :attr:`~isaaclab.sensors.RayCasterData.quat_w`, and -:attr:`~isaaclab.sensors.RayCasterData.ray_hits_w` properties now return ``wp.array`` instead of -``torch.Tensor``. This follows the same pattern as the general warp backend migration described -above. +:attr:`~isaaclab.sensors.RayCasterData.ray_hits_w` properties now return +:class:`~isaaclab.utils.warp.ProxyArray` instead of ``torch.Tensor``. This follows the same +pattern as the general ProxyArray backend migration described above. .. code-block:: python - import warp as wp - # Before (Isaac Lab 2.x) ray_hits = ray_caster.data.ray_hits_w # torch.Tensor sensor_pos = ray_caster.data.pos_w # torch.Tensor # After (Isaac Lab 3.x) - ray_hits = ray_caster.data.ray_hits_w # wp.array - sensor_pos = ray_caster.data.pos_w # wp.array + ray_hits = ray_caster.data.ray_hits_w # ProxyArray + sensor_pos = ray_caster.data.pos_w # ProxyArray - # To use with torch operations, wrap with wp.to_torch() - ray_hits_torch = wp.to_torch(ray_caster.data.ray_hits_w) - sensor_pos_torch = wp.to_torch(ray_caster.data.pos_w) + # To use with torch operations, access .torch + ray_hits_torch = ray_caster.data.ray_hits_w.torch + sensor_pos_torch = ray_caster.data.pos_w.torch Ray Alignment Configuration @@ -1672,6 +1709,60 @@ Deprecated retargeters have been moved to ``isaaclab_teleop.deprecated.openxr.re compatibility. These will be removed in a future release. +.. _torcharray-migration: + +ProxyArray Interop and Temporary Compatibility +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +Asset and sensor data class properties return :class:`~isaaclab.utils.warp.ProxyArray`, a +lightweight wrapper with explicit ``.torch`` and ``.warp`` accessors: + +.. code-block:: python + + # BEFORE (2.x) — properties returned torch.Tensor directly + joint_pos = robot.data.joint_pos # torch.Tensor + root_pos = robot.data.root_pos_w # torch.Tensor + + # AFTER (3.0) — properties return ProxyArray, use .torch for the tensor + joint_pos = robot.data.joint_pos.torch # cached zero-copy torch.Tensor + root_pos = robot.data.root_pos_w.torch # cached zero-copy torch.Tensor + joint_pos_warp = robot.data.joint_pos.warp # the underlying warp.array + +**Automatic interop — in many cases, no changes are needed:** + +- **Warp kernels:** ``ProxyArray`` implements ``__cuda_array_interface__``, so it can be passed + directly to ``wp.launch()`` without calling ``.warp``: + + .. code-block:: python + + # Just works — no .warp needed + wp.launch(my_kernel, inputs=[robot.data.joint_pos], ...) + +- **Torch functions:** ``ProxyArray`` implements ``__torch_function__``, so ``torch.*`` operations + accept it directly. During the deprecation period this emits a one-time warning, but works: + + .. code-block:: python + + # Works (emits DeprecationWarning once, then silent) + mean_pos = torch.mean(robot.data.joint_pos, dim=1) + clipped = torch.clamp(robot.data.joint_pos, -3.14, 3.14) + +**What to change:** + +1. Append ``.torch`` where you need an explicit ``torch.Tensor`` (e.g., for indexing, slicing, + or passing to non-torch libraries). +2. Warp kernel calls need no changes — ``ProxyArray`` works transparently. +3. If you need the underlying ``warp.array`` (e.g., for ``ptr``, ``strides``), use ``.warp``. +4. Replace legacy ``wp.to_torch(proxy_array)`` calls with ``proxy_array.torch``. + +.. note:: + + The ``__torch_function__`` bridge and the temporary ``wp.to_torch(proxy_array)`` shim will + be removed in a future release. We recommend migrating to explicit ``.torch`` access now. + +For a complete guide, see :doc:`/source/how-to/proxy_array`. + + Migration off Deprecated Isaac Sim APIs ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ diff --git a/scripts/benchmarks/benchmark_load_robot.py b/scripts/benchmarks/benchmark_load_robot.py index 1e8dc1f737e9..490a8bd0a064 100644 --- a/scripts/benchmarks/benchmark_load_robot.py +++ b/scripts/benchmarks/benchmark_load_robot.py @@ -57,7 +57,6 @@ imports_time_begin = time.perf_counter_ns() import torch -import warp as wp import isaaclab.sim as sim_utils from isaaclab.assets import ArticulationCfg, AssetBaseCfg @@ -136,15 +135,15 @@ def run_simulator(sim: sim_utils.SimulationContext, scene: InteractiveScene): # root state # we offset the root state by the origin since the states are written in simulation world frame # if this is not done, then the robots will be spawned at the (0, 0, 0) of the simulation world - root_pose = wp.to_torch(robot.data.default_root_pose).clone() + root_pose = robot.data.default_root_pose.torch.clone() root_pose[:, :3] += scene.env_origins robot.write_root_pose_to_sim_index(root_pose=root_pose) - root_vel = wp.to_torch(robot.data.default_root_vel).clone() + root_vel = robot.data.default_root_vel.torch.clone() robot.write_root_velocity_to_sim_index(root_velocity=root_vel) # set joint positions with some noise joint_pos, joint_vel = ( - wp.to_torch(robot.data.default_joint_pos).clone(), - wp.to_torch(robot.data.default_joint_vel).clone(), + robot.data.default_joint_pos.torch.clone(), + robot.data.default_joint_vel.torch.clone(), ) joint_pos += torch.rand_like(joint_pos) * 0.1 robot.write_joint_position_to_sim_index(position=joint_pos) @@ -153,7 +152,7 @@ def run_simulator(sim: sim_utils.SimulationContext, scene: InteractiveScene): scene.reset() # Apply random action # -- generate random joint efforts - efforts = torch.randn_like(wp.to_torch(robot.data.joint_pos)) * 5.0 + efforts = torch.randn_like(robot.data.joint_pos.torch) * 5.0 # -- apply action to the robot robot.set_joint_effort_target_index(target=efforts) # -- write data to sim diff --git a/scripts/demos/arms.py b/scripts/demos/arms.py index a224a44b3703..b08686a8e52c 100644 --- a/scripts/demos/arms.py +++ b/scripts/demos/arms.py @@ -36,7 +36,6 @@ import numpy as np import torch -import warp as wp import isaaclab.sim as sim_utils from isaaclab.assets import Articulation @@ -178,15 +177,15 @@ def run_simulator(sim: sim_utils.SimulationContext, entities: dict[str, Articula # reset the scene entities for index, robot in enumerate(entities.values()): # root state - root_pose = wp.to_torch(robot.data.default_root_pose).clone() + root_pose = robot.data.default_root_pose.torch.clone() root_pose[:, :3] += origins[index] robot.write_root_pose_to_sim_index(root_pose=root_pose) - root_vel = wp.to_torch(robot.data.default_root_vel).clone() + root_vel = robot.data.default_root_vel.torch.clone() robot.write_root_velocity_to_sim_index(root_velocity=root_vel) # set joint positions joint_pos, joint_vel = ( - wp.to_torch(robot.data.default_joint_pos).clone(), - wp.to_torch(robot.data.default_joint_vel).clone(), + robot.data.default_joint_pos.torch.clone(), + robot.data.default_joint_vel.torch.clone(), ) robot.write_joint_position_to_sim_index(position=joint_pos) robot.write_joint_velocity_to_sim_index(velocity=joint_vel) @@ -196,10 +195,8 @@ def run_simulator(sim: sim_utils.SimulationContext, entities: dict[str, Articula # apply random actions to the robots for robot in entities.values(): # generate random joint positions - joint_pos_target = ( - wp.to_torch(robot.data.default_joint_pos) + torch.randn_like(wp.to_torch(robot.data.joint_pos)) * 0.1 - ) - soft_limits = wp.to_torch(robot.data.soft_joint_pos_limits) + joint_pos_target = robot.data.default_joint_pos.torch + torch.randn_like(robot.data.joint_pos.torch) * 0.1 + soft_limits = robot.data.soft_joint_pos_limits.torch joint_pos_target = joint_pos_target.clamp_(soft_limits[..., 0], soft_limits[..., 1]) # apply action to the robot robot.set_joint_position_target_index(target=joint_pos_target) diff --git a/scripts/demos/bin_packing.py b/scripts/demos/bin_packing.py index 05da86396a1c..20d2d656cd86 100644 --- a/scripts/demos/bin_packing.py +++ b/scripts/demos/bin_packing.py @@ -46,7 +46,6 @@ import math import torch -import warp as wp import isaaclab.sim as sim_utils import isaaclab.utils.math as math_utils @@ -267,9 +266,9 @@ def run_simulator(sim: SimulationContext, scene: InteractiveScene) -> None: num_envs = scene.num_envs device = scene.device view_indices = torch.arange(num_envs * num_objects, device=device) - default_pose_w = wp.to_torch(groceries.data.default_body_pose).clone() + default_pose_w = groceries.data.default_body_pose.torch.clone() default_pose_w[..., :3] = default_pose_w[..., :3] + scene.env_origins.unsqueeze(1) - default_vel_w = wp.to_torch(groceries.data.default_body_vel).clone() + default_vel_w = groceries.data.default_body_vel.torch.clone() default_state_w = torch.cat([default_pose_w, default_vel_w], dim=-1) # Define simulation stepping sim_dt = sim.get_physics_dt() diff --git a/scripts/demos/bipeds.py b/scripts/demos/bipeds.py index a8dfdf693962..68f79db14fcb 100644 --- a/scripts/demos/bipeds.py +++ b/scripts/demos/bipeds.py @@ -35,7 +35,6 @@ """Rest everything follows.""" import torch -import warp as wp import isaaclab.sim as sim_utils from isaaclab.assets import Articulation @@ -92,22 +91,22 @@ def run_simulator(sim: sim_utils.SimulationContext, robots: list[Articulation], for index, robot in enumerate(robots): # reset dof state joint_pos, joint_vel = ( - wp.to_torch(robot.data.default_joint_pos), - wp.to_torch(robot.data.default_joint_vel), + robot.data.default_joint_pos.torch, + robot.data.default_joint_vel.torch, ) robot.write_joint_position_to_sim_index(position=joint_pos) robot.write_joint_velocity_to_sim_index(velocity=joint_vel) - root_pose = wp.to_torch(robot.data.default_root_pose).clone() + root_pose = robot.data.default_root_pose.torch.clone() root_pose[:, :3] += origins[index] robot.write_root_pose_to_sim_index(root_pose=root_pose) - root_vel = wp.to_torch(robot.data.default_root_vel).clone() + root_vel = robot.data.default_root_vel.torch.clone() robot.write_root_velocity_to_sim_index(root_velocity=root_vel) robot.reset() # reset command print(">>>>>>>> Reset!") # apply action to the robot for robot in robots: - robot.set_joint_position_target_index(target=wp.to_torch(robot.data.default_joint_pos).clone()) + robot.set_joint_position_target_index(target=robot.data.default_joint_pos.torch.clone()) robot.write_data_to_sim() # perform step sim.step() diff --git a/scripts/demos/deformables.py b/scripts/demos/deformables.py index 4594b6bdea15..71f40b5f9c4a 100644 --- a/scripts/demos/deformables.py +++ b/scripts/demos/deformables.py @@ -38,7 +38,6 @@ import numpy as np import torch import tqdm -import warp as wp # deformables supported in PhysX from isaaclab_physx.assets import DeformableObject, DeformableObjectCfg @@ -207,7 +206,7 @@ def run_simulator(sim: sim_utils.SimulationContext, entities: dict[str, Deformab # reset deformable object state for _, deform_body in enumerate(entities.values()): # root state - nodal_state = wp.to_torch(deform_body.data.default_nodal_state_w).clone() + nodal_state = deform_body.data.default_nodal_state_w.torch.clone() deform_body.write_nodal_state_to_sim_index(nodal_state) # reset the internal state deform_body.reset() diff --git a/scripts/demos/h1_locomotion.py b/scripts/demos/h1_locomotion.py index 2a48828a425d..7e7429c35b0d 100644 --- a/scripts/demos/h1_locomotion.py +++ b/scripts/demos/h1_locomotion.py @@ -45,7 +45,6 @@ """Rest everything follows.""" import torch -import warp as wp from rsl_rl.runners import OnPolicyRunner import carb @@ -199,10 +198,10 @@ def _update_camera(self): """Updates the per-frame transform of the third-person view camera to follow the selected robot's torso transform.""" - base_pos = wp.to_torch(self.env.unwrapped.scene["robot"].data.root_pos_w)[ + base_pos = self.env.unwrapped.scene["robot"].data.root_pos_w.torch[ self._selected_id, : ] # - env.scene.env_origins - base_quat = wp.to_torch(self.env.unwrapped.scene["robot"].data.root_quat_w)[self._selected_id, :] + base_quat = self.env.unwrapped.scene["robot"].data.root_quat_w.torch[self._selected_id, :] camera_pos = quat_apply(base_quat, self._camera_local_transform) + base_pos diff --git a/scripts/demos/hands.py b/scripts/demos/hands.py index 2199d7f5401c..4b9a12d7b433 100644 --- a/scripts/demos/hands.py +++ b/scripts/demos/hands.py @@ -36,7 +36,6 @@ import numpy as np import torch -import warp as wp import isaaclab.sim as sim_utils from isaaclab.assets import Articulation @@ -112,15 +111,15 @@ def run_simulator(sim: sim_utils.SimulationContext, entities: dict[str, Articula # reset robots for index, robot in enumerate(entities.values()): # root state - root_pose = wp.to_torch(robot.data.default_root_pose).clone() + root_pose = robot.data.default_root_pose.torch.clone() root_pose[:, :3] += origins[index] robot.write_root_pose_to_sim_index(root_pose=root_pose) - root_vel = wp.to_torch(robot.data.default_root_vel).clone() + root_vel = robot.data.default_root_vel.torch.clone() robot.write_root_velocity_to_sim_index(root_velocity=root_vel) # joint state joint_pos, joint_vel = ( - wp.to_torch(robot.data.default_joint_pos).clone(), - wp.to_torch(robot.data.default_joint_vel).clone(), + robot.data.default_joint_pos.torch.clone(), + robot.data.default_joint_vel.torch.clone(), ) robot.write_joint_position_to_sim_index(position=joint_pos) robot.write_joint_velocity_to_sim_index(velocity=joint_vel) @@ -133,7 +132,7 @@ def run_simulator(sim: sim_utils.SimulationContext, entities: dict[str, Articula # apply default actions to the hands robots for robot in entities.values(): # generate joint positions - joint_pos_target = wp.to_torch(robot.data.soft_joint_pos_limits)[..., grasp_mode] + joint_pos_target = robot.data.soft_joint_pos_limits.torch[..., grasp_mode] # apply action to the robot robot.set_joint_position_target_index(target=joint_pos_target) # write data to sim diff --git a/scripts/demos/haply_teleoperation.py b/scripts/demos/haply_teleoperation.py index 928d04506e9c..a1c5a7cec677 100644 --- a/scripts/demos/haply_teleoperation.py +++ b/scripts/demos/haply_teleoperation.py @@ -59,7 +59,6 @@ import numpy as np import torch -import warp as wp import isaaclab.sim as sim_utils from isaaclab.assets import Articulation, AssetBaseCfg, RigidObject, RigidObjectCfg @@ -193,9 +192,9 @@ def run_simulator( ee_body_name = "panda_hand" ee_body_idx = robot.body_names.index(ee_body_name) - joint_pos = wp.to_torch(robot.data.default_joint_pos).clone() + joint_pos = robot.data.default_joint_pos.torch.clone() joint_pos[0, :7] = torch.tensor([0.0, -0.569, 0.0, -2.81, 0.0, 3.037, 0.741], device=robot.device) - joint_vel = wp.to_torch(robot.data.default_joint_vel).clone() + joint_vel = robot.data.default_joint_vel.torch.clone() robot.write_joint_position_to_sim_index(position=joint_pos) robot.write_joint_velocity_to_sim_index(velocity=joint_vel) @@ -205,7 +204,7 @@ def run_simulator( scene.update(sim_dt) # Initialize the position of franka - robot_initial_pos = wp.to_torch(robot.data.body_pos_w)[0, ee_body_idx].cpu().numpy() + robot_initial_pos = robot.data.body_pos_w.torch[0, ee_body_idx].cpu().numpy() haply_initial_pos = np.array([0.0, 0.0, 0.0], dtype=np.float32) ik_controller_cfg = DifferentialIKControllerCfg( @@ -228,7 +227,7 @@ def run_simulator( # Initialize IK controller ik_controller = DifferentialIKController(cfg=ik_controller_cfg, num_envs=scene.num_envs, device=sim.device) - initial_ee_quat = wp.to_torch(robot.data.body_quat_w)[:, ee_body_idx] + initial_ee_quat = robot.data.body_quat_w.torch[:, ee_body_idx] ik_controller.set_command(command=torch.zeros(scene.num_envs, 3, device=sim.device), ee_quat=initial_ee_quat) prev_button_a = False @@ -237,7 +236,7 @@ def run_simulator( gripper_target = 0.04 # Initialize the rotation of franka end-effector - ee_rotation_angle = wp.to_torch(robot.data.joint_pos)[0, 6].item() + ee_rotation_angle = robot.data.joint_pos.torch[0, 6].item() rotation_step = np.pi / 3 print("\n[INFO] Teleoperation ready!") @@ -247,22 +246,22 @@ def run_simulator( while simulation_app.is_running(): if count % 10000 == 0: count = 1 - root_pose = wp.to_torch(robot.data.default_root_pose).clone() + root_pose = robot.data.default_root_pose.torch.clone() root_pose[:, :3] += scene.env_origins robot.write_root_pose_to_sim_index(root_pose=root_pose) - root_vel = wp.to_torch(robot.data.default_root_vel).clone() + root_vel = robot.data.default_root_vel.torch.clone() robot.write_root_velocity_to_sim_index(root_velocity=root_vel) - joint_pos = wp.to_torch(robot.data.default_joint_pos).clone() + joint_pos = robot.data.default_joint_pos.torch.clone() joint_pos[0, :7] = torch.tensor([0.0, -0.569, 0.0, -2.81, 0.0, 3.037, 0.741], device=robot.device) - joint_vel = wp.to_torch(robot.data.default_joint_vel).clone() + joint_vel = robot.data.default_joint_vel.torch.clone() robot.write_joint_position_to_sim_index(position=joint_pos) robot.write_joint_velocity_to_sim_index(velocity=joint_vel) - cube_pose = wp.to_torch(cube.data.default_root_pose).clone() + cube_pose = cube.data.default_root_pose.torch.clone() cube_pose[:, :3] += scene.env_origins cube.write_root_pose_to_sim_index(root_pose=cube_pose) - cube_vel = wp.to_torch(cube.data.default_root_vel).clone() + cube_vel = cube.data.default_root_vel.torch.clone() cube.write_root_velocity_to_sim_index(root_velocity=cube_vel) scene.reset() @@ -304,16 +303,16 @@ def run_simulator( target_pos_tensor = torch.tensor(target_pos, dtype=torch.float32, device=sim.device).unsqueeze(0) - current_joint_pos = wp.to_torch(robot.data.joint_pos)[:, arm_joint_indices] - ee_pos_w = wp.to_torch(robot.data.body_pos_w)[:, ee_body_idx] - ee_quat_w = wp.to_torch(robot.data.body_quat_w)[:, ee_body_idx] + current_joint_pos = robot.data.joint_pos.torch[:, arm_joint_indices] + ee_pos_w = robot.data.body_pos_w.torch[:, ee_body_idx] + ee_quat_w = robot.data.body_quat_w.torch[:, ee_body_idx] # get jacobian to IK controller jacobian = robot.root_view.get_jacobians()[:, ee_body_idx, :, arm_joint_indices] ik_controller.set_command(command=target_pos_tensor, ee_quat=ee_quat_w) joint_pos_des = ik_controller.compute(ee_pos_w, ee_quat_w, jacobian, current_joint_pos) - joint_pos_target = wp.to_torch(robot.data.joint_pos)[0].clone() + joint_pos_target = robot.data.joint_pos.torch[0].clone() # Update joints: 6 from IK + 1 from button control (correct by design) joint_pos_target[arm_joint_indices] = joint_pos_des[0] # panda_joint1-6 from IK diff --git a/scripts/demos/multi_asset.py b/scripts/demos/multi_asset.py index 05412f327c39..8091c71afe0e 100644 --- a/scripts/demos/multi_asset.py +++ b/scripts/demos/multi_asset.py @@ -39,8 +39,6 @@ import random -import warp as wp - from pxr import Gf, Sdf import isaaclab.sim as sim_utils @@ -242,28 +240,28 @@ def run_simulator(sim: SimulationContext, scene: InteractiveScene): count = 0 # reset the scene entities # object - root_pose = wp.to_torch(rigid_object.data.default_root_pose).clone() + root_pose = rigid_object.data.default_root_pose.torch.clone() root_pose[:, :3] += scene.env_origins rigid_object.write_root_pose_to_sim_index(root_pose=root_pose) - root_vel = wp.to_torch(rigid_object.data.default_root_vel).clone() + root_vel = rigid_object.data.default_root_vel.torch.clone() rigid_object.write_root_velocity_to_sim_index(root_velocity=root_vel) # object collection - default_pose_w = wp.to_torch(rigid_object_collection.data.default_body_pose).clone() + default_pose_w = rigid_object_collection.data.default_body_pose.torch.clone() default_pose_w[..., :3] += scene.env_origins.unsqueeze(1) rigid_object_collection.write_body_pose_to_sim_index(body_poses=default_pose_w) - default_vel_w = wp.to_torch(rigid_object_collection.data.default_body_vel).clone() + default_vel_w = rigid_object_collection.data.default_body_vel.torch.clone() rigid_object_collection.write_body_com_velocity_to_sim_index(body_velocities=default_vel_w) # robot # -- root state - root_pose = wp.to_torch(robot.data.default_root_pose).clone() + root_pose = robot.data.default_root_pose.torch.clone() root_pose[:, :3] += scene.env_origins robot.write_root_pose_to_sim_index(root_pose=root_pose) - root_vel = wp.to_torch(robot.data.default_root_vel).clone() + root_vel = robot.data.default_root_vel.torch.clone() robot.write_root_velocity_to_sim_index(root_velocity=root_vel) # -- joint state joint_pos, joint_vel = ( - wp.to_torch(robot.data.default_joint_pos).clone(), - wp.to_torch(robot.data.default_joint_vel).clone(), + robot.data.default_joint_pos.torch.clone(), + robot.data.default_joint_vel.torch.clone(), ) robot.write_joint_position_to_sim_index(position=joint_pos) robot.write_joint_velocity_to_sim_index(velocity=joint_vel) @@ -272,7 +270,7 @@ def run_simulator(sim: SimulationContext, scene: InteractiveScene): print("[INFO]: Resetting scene state...") # Apply action to robot - robot.set_joint_position_target_index(target=wp.to_torch(robot.data.default_joint_pos)) + robot.set_joint_position_target_index(target=robot.data.default_joint_pos.torch) # Write data to sim scene.write_data_to_sim() # Perform step diff --git a/scripts/demos/pick_and_place.py b/scripts/demos/pick_and_place.py index af7d54d1cc13..0308706ce38d 100644 --- a/scripts/demos/pick_and_place.py +++ b/scripts/demos/pick_and_place.py @@ -143,8 +143,8 @@ def __init__(self, cfg: PickAndPlaceEnvCfg, render_mode: str | None = None, **kw self._z_dof_idx, _ = self.pick_and_place.find_joints(self.cfg.z_dof_name) # joints info - self.joint_pos = wp.to_torch(self.pick_and_place.data.joint_pos) - self.joint_vel = wp.to_torch(self.pick_and_place.data.joint_vel) + self.joint_pos = self.pick_and_place.data.joint_pos.torch + self.joint_vel = self.pick_and_place.data.joint_vel.torch # Buffers self.go_to_cube = torch.zeros(self.num_envs, dtype=torch.bool, device=self.device) @@ -240,13 +240,13 @@ def _apply_action(self) -> None: # Process each environment independently if self.go_to_cube.any(): # Effort based proportional controller to track the cube position - head_pos_x = wp.to_torch(self.pick_and_place.data.joint_pos)[self.go_to_cube, self._x_dof_idx[0]] - head_pos_y = wp.to_torch(self.pick_and_place.data.joint_pos)[self.go_to_cube, self._y_dof_idx[0]] + head_pos_x = self.pick_and_place.data.joint_pos.torch[self.go_to_cube, self._x_dof_idx[0]] + head_pos_y = self.pick_and_place.data.joint_pos.torch[self.go_to_cube, self._y_dof_idx[0]] cube_pos_x = ( - wp.to_torch(self.cube.data.root_pos_w)[self.go_to_cube, 0] - self.scene.env_origins[self.go_to_cube, 0] + self.cube.data.root_pos_w.torch[self.go_to_cube, 0] - self.scene.env_origins[self.go_to_cube, 0] ) cube_pos_y = ( - wp.to_torch(self.cube.data.root_pos_w)[self.go_to_cube, 1] - self.scene.env_origins[self.go_to_cube, 1] + self.cube.data.root_pos_w.torch[self.go_to_cube, 1] - self.scene.env_origins[self.go_to_cube, 1] ) d_cube_robot_x = cube_pos_x - head_pos_x d_cube_robot_y = cube_pos_y - head_pos_y @@ -255,8 +255,8 @@ def _apply_action(self) -> None: ) if self.go_to_target.any(): # Effort based proportional controller to track the target position - head_pos_x = wp.to_torch(self.pick_and_place.data.joint_pos)[self.go_to_target, self._x_dof_idx[0]] - head_pos_y = wp.to_torch(self.pick_and_place.data.joint_pos)[self.go_to_target, self._y_dof_idx[0]] + head_pos_x = self.pick_and_place.data.joint_pos.torch[self.go_to_target, self._x_dof_idx[0]] + head_pos_y = self.pick_and_place.data.joint_pos.torch[self.go_to_target, self._y_dof_idx[0]] target_pos_x = self.target_pos[self.go_to_target, 0] target_pos_y = self.target_pos[self.go_to_target, 1] d_target_robot_x = target_pos_x - head_pos_x @@ -304,12 +304,12 @@ def _get_rewards(self) -> torch.Tensor: def _get_dones(self) -> tuple[torch.Tensor, torch.Tensor]: # Dones - self.joint_pos = wp.to_torch(self.pick_and_place.data.joint_pos) - self.joint_vel = wp.to_torch(self.pick_and_place.data.joint_vel) + self.joint_pos = self.pick_and_place.data.joint_pos.torch + self.joint_vel = self.pick_and_place.data.joint_vel.torch # Check for time out time_out = self.episode_length_buf >= self.max_episode_length - 1 # Check if the cube reached the target - cube_root_pos_w = wp.to_torch(self.cube.data.root_pos_w) + cube_root_pos_w = self.cube.data.root_pos_w.torch cube_to_target_x_dist = cube_root_pos_w[:, 0] - self.target_pos[:, 0] - self.scene.env_origins[:, 0] cube_to_target_y_dist = cube_root_pos_w[:, 1] - self.target_pos[:, 1] - self.scene.env_origins[:, 1] cube_to_target_z_dist = cube_root_pos_w[:, 2] - self.target_pos[:, 2] - self.scene.env_origins[:, 2] @@ -350,7 +350,7 @@ def _reset_idx(self, env_ids: Sequence[int] | None): self.target_pos[env_ids, 2] = self.cfg.target_z_pos # Set the initial position of the cube - cube_pos = wp.to_torch(self.cube.data.default_root_pose)[env_ids] + cube_pos = self.cube.data.default_root_pose.torch[env_ids] cube_pos[:, 0] = sample_uniform( self.cfg.initial_object_x_pos_range[0], self.cfg.initial_object_x_pos_range[1], @@ -368,7 +368,7 @@ def _reset_idx(self, env_ids: Sequence[int] | None): self.cube.write_root_pose_to_sim_index(root_pose=cube_pos, env_ids=env_ids) # Set the initial position of the robot - joint_pos = wp.to_torch(self.pick_and_place.data.default_joint_pos)[env_ids] + joint_pos = self.pick_and_place.data.default_joint_pos.torch[env_ids] joint_pos[:, self._x_dof_idx] += sample_uniform( self.cfg.initial_x_pos_range[0], self.cfg.initial_x_pos_range[1], @@ -387,7 +387,7 @@ def _reset_idx(self, env_ids: Sequence[int] | None): joint_pos[:, self._z_dof_idx].shape, self.device, ) - joint_vel = wp.to_torch(self.pick_and_place.data.default_joint_vel)[env_ids] + joint_vel = self.pick_and_place.data.default_joint_vel.torch[env_ids] self.joint_pos[env_ids] = joint_pos self.joint_vel[env_ids] = joint_vel diff --git a/scripts/demos/quadcopter.py b/scripts/demos/quadcopter.py index 0e665596e71d..4cc9e1a27095 100644 --- a/scripts/demos/quadcopter.py +++ b/scripts/demos/quadcopter.py @@ -35,7 +35,6 @@ """Rest everything follows.""" import torch -import warp as wp import isaaclab.sim as sim_utils from isaaclab.assets import Articulation @@ -93,12 +92,12 @@ def main(): sim_time = 0.0 count = 0 # reset dof state - joint_pos, joint_vel = wp.to_torch(robot.data.default_joint_pos), wp.to_torch(robot.data.default_joint_vel) + joint_pos, joint_vel = robot.data.default_joint_pos.torch, robot.data.default_joint_vel.torch robot.write_joint_position_to_sim_index(position=joint_pos) robot.write_joint_velocity_to_sim_index(velocity=joint_vel) - default_root_pose = wp.to_torch(robot.data.default_root_pose) + default_root_pose = robot.data.default_root_pose.torch robot.write_root_pose_to_sim_index(root_pose=default_root_pose) - default_root_vel = wp.to_torch(robot.data.default_root_vel) + default_root_vel = robot.data.default_root_vel.torch robot.write_root_velocity_to_sim_index(root_velocity=default_root_vel) robot.reset() # reset command diff --git a/scripts/demos/quadrupeds.py b/scripts/demos/quadrupeds.py index 28cc64074bff..59551a3738a2 100644 --- a/scripts/demos/quadrupeds.py +++ b/scripts/demos/quadrupeds.py @@ -36,7 +36,6 @@ import numpy as np import torch -import warp as wp import isaaclab.sim as sim_utils from isaaclab.assets import Articulation @@ -141,15 +140,15 @@ def run_simulator(sim: sim_utils.SimulationContext, entities: dict[str, Articula # reset robots for index, robot in enumerate(entities.values()): # root state - root_pose = wp.to_torch(robot.data.default_root_pose).clone() + root_pose = robot.data.default_root_pose.torch.clone() root_pose[:, :3] += origins[index] robot.write_root_pose_to_sim_index(root_pose=root_pose) - root_vel = wp.to_torch(robot.data.default_root_vel).clone() + root_vel = robot.data.default_root_vel.torch.clone() robot.write_root_velocity_to_sim_index(root_velocity=root_vel) # joint state joint_pos, joint_vel = ( - wp.to_torch(robot.data.default_joint_pos).clone(), - wp.to_torch(robot.data.default_joint_vel).clone(), + robot.data.default_joint_pos.torch.clone(), + robot.data.default_joint_vel.torch.clone(), ) robot.write_joint_position_to_sim_index(position=joint_pos) robot.write_joint_velocity_to_sim_index(velocity=joint_vel) @@ -159,9 +158,7 @@ def run_simulator(sim: sim_utils.SimulationContext, entities: dict[str, Articula # apply default actions to the quadrupedal robots for robot in entities.values(): # generate random joint positions - joint_pos_target = ( - wp.to_torch(robot.data.default_joint_pos) + torch.randn_like(wp.to_torch(robot.data.joint_pos)) * 0.1 - ) + joint_pos_target = robot.data.default_joint_pos.torch + torch.randn_like(robot.data.joint_pos.torch) * 0.1 # apply action to the robot robot.set_joint_position_target_index(target=joint_pos_target) # write data to sim diff --git a/scripts/demos/sensors/cameras.py b/scripts/demos/sensors/cameras.py index f96f7226039f..9415c2a5a66e 100644 --- a/scripts/demos/sensors/cameras.py +++ b/scripts/demos/sensors/cameras.py @@ -44,7 +44,6 @@ import matplotlib.pyplot as plt import numpy as np import torch -import warp as wp import isaaclab.sim as sim_utils from isaaclab.assets import ArticulationCfg, AssetBaseCfg @@ -192,15 +191,15 @@ def run_simulator(sim: sim_utils.SimulationContext, scene: InteractiveScene): # root state # we offset the root state by the origin since the states are written in simulation world frame # if this is not done, then the robots will be spawned at the (0, 0, 0) of the simulation world - root_pose = wp.to_torch(scene["robot"].data.default_root_pose).clone() + root_pose = scene["robot"].data.default_root_pose.torch.clone() root_pose[:, :3] += scene.env_origins scene["robot"].write_root_pose_to_sim_index(root_pose=root_pose) - root_vel = wp.to_torch(scene["robot"].data.default_root_vel).clone() + root_vel = scene["robot"].data.default_root_vel.torch.clone() scene["robot"].write_root_velocity_to_sim_index(root_velocity=root_vel) # set joint positions with some noise joint_pos, joint_vel = ( - wp.to_torch(scene["robot"].data.default_joint_pos).clone(), - wp.to_torch(scene["robot"].data.default_joint_vel).clone(), + scene["robot"].data.default_joint_pos.torch.clone(), + scene["robot"].data.default_joint_vel.torch.clone(), ) joint_pos += torch.rand_like(joint_pos) * 0.1 scene["robot"].write_joint_position_to_sim_index(position=joint_pos) @@ -210,7 +209,7 @@ def run_simulator(sim: sim_utils.SimulationContext, scene: InteractiveScene): print("[INFO]: Resetting robot state...") # Apply default actions to the robot # -- generate actions/commands - targets = wp.to_torch(scene["robot"].data.default_joint_pos) + targets = scene["robot"].data.default_joint_pos.torch # -- apply action to the robot scene["robot"].set_joint_position_target_index(target=targets) # -- write data to sim diff --git a/scripts/demos/sensors/contact_sensor.py b/scripts/demos/sensors/contact_sensor.py index 0c1a72126e90..37ddd14e28a6 100644 --- a/scripts/demos/sensors/contact_sensor.py +++ b/scripts/demos/sensors/contact_sensor.py @@ -26,7 +26,6 @@ """Rest everything follows.""" import torch -import warp as wp import isaaclab.sim as sim_utils from isaaclab.assets import AssetBaseCfg, RigidObjectCfg @@ -109,15 +108,15 @@ def run_simulator(sim: sim_utils.SimulationContext, scene: InteractiveScene): # root state # we offset the root state by the origin since the states are written in simulation world frame # if this is not done, then the robots will be spawned at the (0, 0, 0) of the simulation world - root_pose = wp.to_torch(scene["robot"].data.default_root_pose).clone() + root_pose = scene["robot"].data.default_root_pose.torch.clone() root_pose[:, :3] += scene.env_origins scene["robot"].write_root_pose_to_sim_index(root_pose=root_pose) - root_vel = wp.to_torch(scene["robot"].data.default_root_vel).clone() + root_vel = scene["robot"].data.default_root_vel.torch.clone() scene["robot"].write_root_velocity_to_sim_index(root_velocity=root_vel) # set joint positions with some noise joint_pos, joint_vel = ( - wp.to_torch(scene["robot"].data.default_joint_pos).clone(), - wp.to_torch(scene["robot"].data.default_joint_vel).clone(), + scene["robot"].data.default_joint_pos.torch.clone(), + scene["robot"].data.default_joint_vel.torch.clone(), ) joint_pos += torch.rand_like(joint_pos) * 0.1 scene["robot"].write_joint_position_to_sim_index(position=joint_pos) @@ -127,7 +126,7 @@ def run_simulator(sim: sim_utils.SimulationContext, scene: InteractiveScene): print("[INFO]: Resetting robot state...") # Apply default actions to the robot # -- generate actions/commands - targets = wp.to_torch(scene["robot"].data.default_joint_pos) + targets = scene["robot"].data.default_joint_pos.torch # -- apply action to the robot scene["robot"].set_joint_position_target_index(target=targets) # -- write data to sim diff --git a/scripts/demos/sensors/frame_transformer_sensor.py b/scripts/demos/sensors/frame_transformer_sensor.py index e8a8760219a9..039e38935547 100644 --- a/scripts/demos/sensors/frame_transformer_sensor.py +++ b/scripts/demos/sensors/frame_transformer_sensor.py @@ -24,7 +24,6 @@ """Rest everything follows.""" import torch -import warp as wp import isaaclab.sim as sim_utils from isaaclab.assets import AssetBaseCfg, RigidObjectCfg @@ -105,15 +104,15 @@ def run_simulator(sim: sim_utils.SimulationContext, scene: InteractiveScene): # root state # we offset the root state by the origin since the states are written in simulation world frame # if this is not done, then the robots will be spawned at the (0, 0, 0) of the simulation world - root_pose = wp.to_torch(scene["robot"].data.default_root_pose).clone() + root_pose = scene["robot"].data.default_root_pose.torch.clone() root_pose[:, :3] += scene.env_origins scene["robot"].write_root_pose_to_sim_index(root_pose=root_pose) - root_vel = wp.to_torch(scene["robot"].data.default_root_vel).clone() + root_vel = scene["robot"].data.default_root_vel.torch.clone() scene["robot"].write_root_velocity_to_sim_index(root_velocity=root_vel) # set joint positions with some noise joint_pos, joint_vel = ( - wp.to_torch(scene["robot"].data.default_joint_pos).clone(), - wp.to_torch(scene["robot"].data.default_joint_vel).clone(), + scene["robot"].data.default_joint_pos.torch.clone(), + scene["robot"].data.default_joint_vel.torch.clone(), ) joint_pos += torch.rand_like(joint_pos) * 0.1 scene["robot"].write_joint_position_to_sim_index(position=joint_pos) @@ -123,7 +122,7 @@ def run_simulator(sim: sim_utils.SimulationContext, scene: InteractiveScene): print("[INFO]: Resetting robot state...") # Apply default actions to the robot # -- generate actions/commands - targets = wp.to_torch(scene["robot"].data.default_joint_pos) + targets = scene["robot"].data.default_joint_pos.torch # -- apply action to the robot scene["robot"].set_joint_position_target_index(target=targets) # -- write data to sim diff --git a/scripts/demos/sensors/imu_sensor.py b/scripts/demos/sensors/imu_sensor.py index f5c79529be81..beb701d62074 100644 --- a/scripts/demos/sensors/imu_sensor.py +++ b/scripts/demos/sensors/imu_sensor.py @@ -26,7 +26,6 @@ """Rest everything follows.""" import torch -import warp as wp import isaaclab.sim as sim_utils from isaaclab.assets import AssetBaseCfg @@ -73,15 +72,15 @@ def run_simulator(sim: sim_utils.SimulationContext, scene: InteractiveScene): # reset counter count = 0 # reset the scene entities - root_pose = wp.to_torch(scene["robot"].data.default_root_pose).clone() + root_pose = scene["robot"].data.default_root_pose.torch.clone() root_pose[:, :3] += scene.env_origins scene["robot"].write_root_link_pose_to_sim_index(root_pose=root_pose) - root_vel = wp.to_torch(scene["robot"].data.default_root_vel).clone() + root_vel = scene["robot"].data.default_root_vel.torch.clone() scene["robot"].write_root_com_velocity_to_sim_index(root_velocity=root_vel) # set joint positions with some noise joint_pos, joint_vel = ( - wp.to_torch(scene["robot"].data.default_joint_pos).clone(), - wp.to_torch(scene["robot"].data.default_joint_vel).clone(), + scene["robot"].data.default_joint_pos.torch.clone(), + scene["robot"].data.default_joint_vel.torch.clone(), ) joint_pos += torch.rand_like(joint_pos) * 0.1 scene["robot"].write_joint_position_to_sim_index(position=joint_pos) @@ -90,7 +89,7 @@ def run_simulator(sim: sim_utils.SimulationContext, scene: InteractiveScene): scene.reset() print("[INFO]: Resetting robot state...") # Apply default actions to the robot - targets = wp.to_torch(scene["robot"].data.default_joint_pos) + targets = scene["robot"].data.default_joint_pos.torch scene["robot"].set_joint_position_target_index(target=targets) scene.write_data_to_sim() # perform step diff --git a/scripts/demos/sensors/multi_mesh_raycaster.py b/scripts/demos/sensors/multi_mesh_raycaster.py index 1a586a8aa2d1..c154a6181763 100644 --- a/scripts/demos/sensors/multi_mesh_raycaster.py +++ b/scripts/demos/sensors/multi_mesh_raycaster.py @@ -49,7 +49,6 @@ import random import torch -import warp as wp from pxr import Gf, Sdf @@ -239,17 +238,17 @@ def run_simulator(sim: sim_utils.SimulationContext, scene: InteractiveScene): count = 0 # reset the scene entities # root state - root_pose = wp.to_torch(scene["asset"].data.default_root_pose).clone() + root_pose = scene["asset"].data.default_root_pose.torch.clone() root_pose[:, :3] += scene.env_origins scene["asset"].write_root_pose_to_sim_index(root_pose=root_pose) - root_vel = wp.to_torch(scene["asset"].data.default_root_vel).clone() + root_vel = scene["asset"].data.default_root_vel.torch.clone() scene["asset"].write_root_velocity_to_sim_index(root_velocity=root_vel) if isinstance(scene["asset"], Articulation): # set joint positions with some noise joint_pos, joint_vel = ( - wp.to_torch(scene["asset"].data.default_joint_pos).clone(), - wp.to_torch(scene["asset"].data.default_joint_vel).clone(), + scene["asset"].data.default_joint_pos.torch.clone(), + scene["asset"].data.default_joint_vel.torch.clone(), ) joint_pos += torch.rand_like(joint_pos) * 0.1 scene["asset"].write_joint_position_to_sim_index(position=joint_pos) @@ -260,7 +259,7 @@ def run_simulator(sim: sim_utils.SimulationContext, scene: InteractiveScene): if isinstance(scene["asset"], Articulation): # -- generate actions/commands - default_joint_pos = wp.to_torch(scene["asset"].data.default_joint_pos) + default_joint_pos = scene["asset"].data.default_joint_pos.torch targets = default_joint_pos + 5 * (torch.rand_like(default_joint_pos) - 0.5) # -- apply action to the asset scene["asset"].set_joint_position_target_index(target=targets) diff --git a/scripts/demos/sensors/multi_mesh_raycaster_camera.py b/scripts/demos/sensors/multi_mesh_raycaster_camera.py index 6c61a2a4c4dd..5d1dfd39a07d 100644 --- a/scripts/demos/sensors/multi_mesh_raycaster_camera.py +++ b/scripts/demos/sensors/multi_mesh_raycaster_camera.py @@ -49,7 +49,6 @@ import random import torch -import warp as wp from pxr import Gf, Sdf @@ -255,17 +254,17 @@ def run_simulator(sim: sim_utils.SimulationContext, scene: InteractiveScene): count = 0 # reset the scene entities # root state - root_pose = wp.to_torch(scene["asset"].data.default_root_pose).clone() + root_pose = scene["asset"].data.default_root_pose.torch.clone() root_pose[:, :3] += scene.env_origins scene["asset"].write_root_pose_to_sim_index(root_pose=root_pose) - root_vel = wp.to_torch(scene["asset"].data.default_root_vel).clone() + root_vel = scene["asset"].data.default_root_vel.torch.clone() scene["asset"].write_root_velocity_to_sim_index(root_velocity=root_vel) if isinstance(scene["asset"], Articulation): # set joint positions with some noise joint_pos, joint_vel = ( - wp.to_torch(scene["asset"].data.default_joint_pos).clone(), - wp.to_torch(scene["asset"].data.default_joint_vel).clone(), + scene["asset"].data.default_joint_pos.torch.clone(), + scene["asset"].data.default_joint_vel.torch.clone(), ) joint_pos += torch.rand_like(joint_pos) * 0.1 scene["asset"].write_joint_position_to_sim_index(position=joint_pos) @@ -276,7 +275,7 @@ def run_simulator(sim: sim_utils.SimulationContext, scene: InteractiveScene): if isinstance(scene["asset"], Articulation): # -- generate actions/commands - default_joint_pos = wp.to_torch(scene["asset"].data.default_joint_pos) + default_joint_pos = scene["asset"].data.default_joint_pos.torch targets = default_joint_pos + 5 * (torch.rand_like(default_joint_pos) - 0.5) # -- apply action to the asset scene["asset"].set_joint_position_target_index(target=targets) @@ -295,7 +294,7 @@ def run_simulator(sim: sim_utils.SimulationContext, scene: InteractiveScene): countdown -= 1 continue - data = scene["ray_caster"].data.ray_hits_w.cpu().numpy() # noqa: F841 + data = scene["ray_caster"].data.ray_hits_w.torch.cpu().numpy() # noqa: F841 triggered = True else: continue diff --git a/scripts/demos/sensors/pva_sensor.py b/scripts/demos/sensors/pva_sensor.py index 7d00056defef..a51ea7c17cc9 100644 --- a/scripts/demos/sensors/pva_sensor.py +++ b/scripts/demos/sensors/pva_sensor.py @@ -26,7 +26,6 @@ """Rest everything follows.""" import torch -import warp as wp import isaaclab.sim as sim_utils from isaaclab.assets import AssetBaseCfg @@ -76,15 +75,15 @@ def run_simulator(sim: sim_utils.SimulationContext, scene: InteractiveScene): # root state # we offset the root state by the origin since the states are written in simulation world frame # if this is not done, then the robots will be spawned at the (0, 0, 0) of the simulation world - root_pose = wp.to_torch(scene["robot"].data.default_root_pose).clone() + root_pose = scene["robot"].data.default_root_pose.torch.clone() root_pose[:, :3] += scene.env_origins scene["robot"].write_root_link_pose_to_sim_index(root_pose=root_pose) - root_vel = wp.to_torch(scene["robot"].data.default_root_vel).clone() + root_vel = scene["robot"].data.default_root_vel.torch.clone() scene["robot"].write_root_com_velocity_to_sim_index(root_velocity=root_vel) # set joint positions with some noise joint_pos, joint_vel = ( - wp.to_torch(scene["robot"].data.default_joint_pos).clone(), - wp.to_torch(scene["robot"].data.default_joint_vel).clone(), + scene["robot"].data.default_joint_pos.torch.clone(), + scene["robot"].data.default_joint_vel.torch.clone(), ) joint_pos += torch.rand_like(joint_pos) * 0.1 scene["robot"].write_joint_position_to_sim_index(position=joint_pos) @@ -94,7 +93,7 @@ def run_simulator(sim: sim_utils.SimulationContext, scene: InteractiveScene): print("[INFO]: Resetting robot state...") # Apply default actions to the robot # -- generate actions/commands - targets = wp.to_torch(scene["robot"].data.default_joint_pos) + targets = scene["robot"].data.default_joint_pos.torch # -- apply action to the robot scene["robot"].set_joint_position_target_index(target=targets) # -- write data to sim diff --git a/scripts/demos/sensors/raycaster_sensor.py b/scripts/demos/sensors/raycaster_sensor.py index dd0b454ad636..25073af40487 100644 --- a/scripts/demos/sensors/raycaster_sensor.py +++ b/scripts/demos/sensors/raycaster_sensor.py @@ -25,7 +25,6 @@ import numpy as np import torch -import warp as wp import isaaclab.sim as sim_utils from isaaclab.assets import AssetBaseCfg @@ -93,15 +92,15 @@ def run_simulator(sim: sim_utils.SimulationContext, scene: InteractiveScene): # root state # we offset the root state by the origin since the states are written in simulation world frame # if this is not done, then the robots will be spawned at the (0, 0, 0) of the simulation world - root_pose = wp.to_torch(scene["robot"].data.default_root_pose).clone() + root_pose = scene["robot"].data.default_root_pose.torch.clone() root_pose[:, :3] += scene.env_origins scene["robot"].write_root_pose_to_sim_index(root_pose=root_pose) - root_vel = wp.to_torch(scene["robot"].data.default_root_vel).clone() + root_vel = scene["robot"].data.default_root_vel.torch.clone() scene["robot"].write_root_velocity_to_sim_index(root_velocity=root_vel) # set joint positions with some noise joint_pos, joint_vel = ( - wp.to_torch(scene["robot"].data.default_joint_pos).clone(), - wp.to_torch(scene["robot"].data.default_joint_vel).clone(), + scene["robot"].data.default_joint_pos.torch.clone(), + scene["robot"].data.default_joint_vel.torch.clone(), ) joint_pos += torch.rand_like(joint_pos) * 0.1 scene["robot"].write_joint_position_to_sim_index(position=joint_pos) @@ -111,7 +110,7 @@ def run_simulator(sim: sim_utils.SimulationContext, scene: InteractiveScene): print("[INFO]: Resetting robot state...") # Apply default actions to the robot # -- generate actions/commands - targets = wp.to_torch(scene["robot"].data.default_joint_pos) + targets = scene["robot"].data.default_joint_pos.torch # -- apply action to the robot scene["robot"].set_joint_position_target_index(target=targets) # -- write data to sim @@ -127,13 +126,13 @@ def run_simulator(sim: sim_utils.SimulationContext, scene: InteractiveScene): # print information from the sensors print("-------------------------------") print(scene["ray_caster"]) - print("Ray cast hit results: ", wp.to_torch(scene["ray_caster"].data.ray_hits_w)) + print("Ray cast hit results: ", scene["ray_caster"].data.ray_hits_w.torch) if not triggered: if countdown > 0: countdown -= 1 continue - data = wp.to_torch(scene["ray_caster"].data.ray_hits_w).cpu().numpy() + data = scene["ray_caster"].data.ray_hits_w.torch.cpu().numpy() np.save("cast_data.npy", data) triggered = True else: diff --git a/scripts/demos/sensors/tacsl_sensor.py b/scripts/demos/sensors/tacsl_sensor.py index 97f59589fa3d..63a215336e94 100644 --- a/scripts/demos/sensors/tacsl_sensor.py +++ b/scripts/demos/sensors/tacsl_sensor.py @@ -30,7 +30,6 @@ import cv2 import numpy as np import torch -import warp as wp from isaaclab.app import AppLauncher @@ -333,10 +332,10 @@ def run_simulator(sim: sim_utils.SimulationContext, scene: InteractiveScene): # Reset robot and contact object positions count = 0 for entity in entity_list: - root_pose = wp.to_torch(scene[entity].data.default_root_pose).clone() + root_pose = scene[entity].data.default_root_pose.torch.clone() root_pose[:, :3] += scene.env_origins scene[entity].write_root_pose_to_sim_index(root_pose=root_pose) - root_vel = wp.to_torch(scene[entity].data.default_root_vel).clone() + root_vel = scene[entity].data.default_root_vel.torch.clone() scene[entity].write_root_velocity_to_sim_index(root_velocity=root_vel) scene.reset() diff --git a/scripts/environments/state_machine/lift_cube_sm.py b/scripts/environments/state_machine/lift_cube_sm.py index 003304089055..ff7734a69e1a 100644 --- a/scripts/environments/state_machine/lift_cube_sm.py +++ b/scripts/environments/state_machine/lift_cube_sm.py @@ -284,12 +284,12 @@ def main(): # -- end-effector frame ee_frame_sensor = env.unwrapped.scene["ee_frame"] tcp_rest_position = ( - wp.to_torch(ee_frame_sensor.data.target_pos_w)[..., 0, :].clone() - env.unwrapped.scene.env_origins + ee_frame_sensor.data.target_pos_w.torch[..., 0, :].clone() - env.unwrapped.scene.env_origins ) - tcp_rest_orientation = wp.to_torch(ee_frame_sensor.data.target_quat_w)[..., 0, :].clone() + tcp_rest_orientation = ee_frame_sensor.data.target_quat_w.torch[..., 0, :].clone() # -- object frame object_data: RigidObjectData = env.unwrapped.scene["object"].data - object_position = wp.to_torch(object_data.root_pos_w) - env.unwrapped.scene.env_origins + object_position = object_data.root_pos_w.torch - env.unwrapped.scene.env_origins # -- target object frame desired_position = env.unwrapped.command_manager.get_command("object_pose")[..., :3] diff --git a/scripts/environments/state_machine/lift_teddy_bear.py b/scripts/environments/state_machine/lift_teddy_bear.py index dd6574049dbe..0aa8e0c33504 100644 --- a/scripts/environments/state_machine/lift_teddy_bear.py +++ b/scripts/environments/state_machine/lift_teddy_bear.py @@ -304,12 +304,12 @@ def main(): # -- end-effector frame ee_frame_sensor = env.unwrapped.scene["ee_frame"] tcp_rest_position = ( - wp.to_torch(ee_frame_sensor.data.target_pos_w)[..., 0, :].clone() - env.unwrapped.scene.env_origins + ee_frame_sensor.data.target_pos_w.torch[..., 0, :].clone() - env.unwrapped.scene.env_origins ) - tcp_rest_orientation = wp.to_torch(ee_frame_sensor.data.target_quat_w)[..., 0, :].clone() + tcp_rest_orientation = ee_frame_sensor.data.target_quat_w.torch[..., 0, :].clone() # -- object frame object_data: RigidObjectData = env.unwrapped.scene["object"].data - object_position = wp.to_torch(object_data.root_pos_w) - env.unwrapped.scene.env_origins + object_position = object_data.root_pos_w.torch - env.unwrapped.scene.env_origins object_position += object_local_grasp_position # -- target object frame diff --git a/scripts/environments/state_machine/open_cabinet_sm.py b/scripts/environments/state_machine/open_cabinet_sm.py index 4b2d52ebe243..39cae64bce07 100644 --- a/scripts/environments/state_machine/open_cabinet_sm.py +++ b/scripts/environments/state_machine/open_cabinet_sm.py @@ -300,16 +300,14 @@ def main(): # observations # -- end-effector frame ee_frame_tf: FrameTransformer = env.unwrapped.scene["ee_frame"] - tcp_rest_position = ( - wp.to_torch(ee_frame_tf.data.target_pos_w)[..., 0, :].clone() - env.unwrapped.scene.env_origins - ) - tcp_rest_orientation = wp.to_torch(ee_frame_tf.data.target_quat_w)[..., 0, :].clone() + tcp_rest_position = ee_frame_tf.data.target_pos_w.torch[..., 0, :].clone() - env.unwrapped.scene.env_origins + tcp_rest_orientation = ee_frame_tf.data.target_quat_w.torch[..., 0, :].clone() # -- handle frame cabinet_frame_tf: FrameTransformer = env.unwrapped.scene["cabinet_frame"] cabinet_position = ( - wp.to_torch(cabinet_frame_tf.data.target_pos_w)[..., 0, :].clone() - env.unwrapped.scene.env_origins + cabinet_frame_tf.data.target_pos_w.torch[..., 0, :].clone() - env.unwrapped.scene.env_origins ) - cabinet_orientation = wp.to_torch(cabinet_frame_tf.data.target_quat_w)[..., 0, :].clone() + cabinet_orientation = cabinet_frame_tf.data.target_quat_w.torch[..., 0, :].clone() # advance state machine actions = open_sm.compute( diff --git a/scripts/imitation_learning/locomanipulation_sdg/generate_data.py b/scripts/imitation_learning/locomanipulation_sdg/generate_data.py index 2b3af68914c2..bbae0689ff65 100644 --- a/scripts/imitation_learning/locomanipulation_sdg/generate_data.py +++ b/scripts/imitation_learning/locomanipulation_sdg/generate_data.py @@ -350,26 +350,26 @@ def project_robot_state_into_env(env: LocomanipulationSDGEnv, input_episode_data root_velocity=torch.zeros((1, 6), device=env.device), env_ids=[0] ) # Update default root pose and velocity for correct state on reset - default_pose = wp.to_torch(env.scene["robot"].data.default_root_pose).clone() + default_pose = env.scene["robot"].data.default_root_pose.torch.clone() default_pose[0] = new_robot_pose[0] - env.scene["robot"].data.default_root_pose.assign( + env.scene["robot"].data.default_root_pose.warp.assign( wp.from_torch(default_pose.to(env.device).contiguous()).view(wp.transformf) ) - default_vel = wp.to_torch(env.scene["robot"].data.default_root_vel).clone() + default_vel = env.scene["robot"].data.default_root_vel.torch.clone() default_vel[0] = torch.zeros(6, device=env.device) - env.scene["robot"].data.default_root_vel.assign(wp.from_torch(default_vel.to(env.device).contiguous())) + env.scene["robot"].data.default_root_vel.warp.assign(wp.from_torch(default_vel.to(env.device).contiguous())) robot_state = recording_initial_state["articulation"]["robot"] joint_position = robot_state["joint_position"][0].to(env.device) joint_velocity = robot_state["joint_velocity"][0].to(env.device) # Update default joint positions and velocities for correct state on reset - default_joint_pos = wp.to_torch(env.scene["robot"].data.default_joint_pos).clone() + default_joint_pos = env.scene["robot"].data.default_joint_pos.torch.clone() default_joint_pos[0] = joint_position - env.scene["robot"].data.default_joint_pos.assign(wp.from_torch(default_joint_pos.to(env.device).contiguous())) - default_joint_vel = wp.to_torch(env.scene["robot"].data.default_joint_vel).clone() + env.scene["robot"].data.default_joint_pos.warp.assign(wp.from_torch(default_joint_pos.to(env.device).contiguous())) + default_joint_vel = env.scene["robot"].data.default_joint_vel.torch.clone() default_joint_vel[0] = joint_velocity - env.scene["robot"].data.default_joint_vel.assign(wp.from_torch(default_joint_vel.to(env.device).contiguous())) + env.scene["robot"].data.default_joint_vel.warp.assign(wp.from_torch(default_joint_vel.to(env.device).contiguous())) env.scene["robot"].write_joint_position_to_sim_index(position=joint_position[None, :], env_ids=[0]) env.scene["robot"].write_joint_velocity_to_sim_index(velocity=joint_velocity[None, :], env_ids=[0]) @@ -401,14 +401,14 @@ def project_object_state_into_env(env: LocomanipulationSDGEnv, input_episode_dat root_velocity=torch.zeros((1, 6), device=env.device), env_ids=[0] ) # Update default root pose and velocity for correct state on reset - default_pose = wp.to_torch(env.scene["object"].data.default_root_pose).clone() + default_pose = env.scene["object"].data.default_root_pose.torch.clone() default_pose[0] = new_object_pose[0] - env.scene["object"].data.default_root_pose.assign( + env.scene["object"].data.default_root_pose.warp.assign( wp.from_torch(default_pose.to(env.device).contiguous()).view(wp.transformf) ) - default_vel = wp.to_torch(env.scene["object"].data.default_root_vel).clone() + default_vel = env.scene["object"].data.default_root_vel.torch.clone() default_vel[0] = torch.zeros(6, device=env.device) - env.scene["object"].data.default_root_vel.assign(wp.from_torch(default_vel.to(env.device).contiguous())) + env.scene["object"].data.default_root_vel.warp.assign(wp.from_torch(default_vel.to(env.device).contiguous())) return new_object_pose diff --git a/scripts/tutorials/01_assets/add_new_robot.py b/scripts/tutorials/01_assets/add_new_robot.py index 087aa865a84c..cbc1e9ab5f03 100644 --- a/scripts/tutorials/01_assets/add_new_robot.py +++ b/scripts/tutorials/01_assets/add_new_robot.py @@ -23,7 +23,6 @@ import numpy as np import torch -import warp as wp import isaaclab.sim as sim_utils from isaaclab.actuators import ImplicitActuatorCfg @@ -110,29 +109,29 @@ def run_simulator(sim: sim_utils.SimulationContext, scene: InteractiveScene): # reset counters count = 0 # reset the scene entities to their initial positions offset by the environment origins - root_jetbot_pose = wp.to_torch(scene["Jetbot"].data.default_root_pose).clone() + root_jetbot_pose = scene["Jetbot"].data.default_root_pose.torch.clone() root_jetbot_pose[:, :3] += scene.env_origins - root_dofbot_pose = wp.to_torch(scene["Dofbot"].data.default_root_pose).clone() + root_dofbot_pose = scene["Dofbot"].data.default_root_pose.torch.clone() root_dofbot_pose[:, :3] += scene.env_origins # copy the default root state to the sim for the jetbot's orientation and velocity scene["Jetbot"].write_root_pose_to_sim_index(root_pose=root_jetbot_pose) - root_jetbot_vel = wp.to_torch(scene["Jetbot"].data.default_root_vel).clone() + root_jetbot_vel = scene["Jetbot"].data.default_root_vel.torch.clone() scene["Jetbot"].write_root_velocity_to_sim_index(root_velocity=root_jetbot_vel) scene["Dofbot"].write_root_pose_to_sim_index(root_pose=root_dofbot_pose) - root_dofbot_vel = wp.to_torch(scene["Dofbot"].data.default_root_vel).clone() + root_dofbot_vel = scene["Dofbot"].data.default_root_vel.torch.clone() scene["Dofbot"].write_root_velocity_to_sim_index(root_velocity=root_dofbot_vel) # copy the default joint states to the sim joint_pos, joint_vel = ( - wp.to_torch(scene["Jetbot"].data.default_joint_pos).clone(), - wp.to_torch(scene["Jetbot"].data.default_joint_vel).clone(), + scene["Jetbot"].data.default_joint_pos.torch.clone(), + scene["Jetbot"].data.default_joint_vel.torch.clone(), ) scene["Jetbot"].write_joint_position_to_sim_index(position=joint_pos) scene["Jetbot"].write_joint_velocity_to_sim_index(velocity=joint_vel) joint_pos, joint_vel = ( - wp.to_torch(scene["Dofbot"].data.default_joint_pos).clone(), - wp.to_torch(scene["Dofbot"].data.default_joint_vel).clone(), + scene["Dofbot"].data.default_joint_pos.torch.clone(), + scene["Dofbot"].data.default_joint_vel.torch.clone(), ) scene["Dofbot"].write_joint_position_to_sim_index(position=joint_pos) scene["Dofbot"].write_joint_velocity_to_sim_index(velocity=joint_vel) @@ -151,7 +150,7 @@ def run_simulator(sim: sim_utils.SimulationContext, scene: InteractiveScene): scene["Jetbot"].set_joint_velocity_target(action) # wave - wave_action = wp.to_torch(scene["Dofbot"].data.default_joint_pos) + wave_action = scene["Dofbot"].data.default_joint_pos.torch wave_action[:, 0:4] = 0.25 * np.sin(2 * np.pi * 0.5 * sim_time) scene["Dofbot"].set_joint_position_target_index(target=wave_action) diff --git a/scripts/tutorials/01_assets/run_articulation.py b/scripts/tutorials/01_assets/run_articulation.py index b6413b0be006..944edf943d4e 100644 --- a/scripts/tutorials/01_assets/run_articulation.py +++ b/scripts/tutorials/01_assets/run_articulation.py @@ -33,7 +33,6 @@ """Rest everything follows.""" import torch -import warp as wp import isaaclab.sim as sim_utils from isaaclab.assets import Articulation @@ -91,15 +90,15 @@ def run_simulator(sim: sim_utils.SimulationContext, entities: dict[str, Articula # root state # we offset the root state by the origin since the states are written in simulation world frame # if this is not done, then the robots will be spawned at the (0, 0, 0) of the simulation world - root_pose = wp.to_torch(robot.data.default_root_pose).clone() + root_pose = robot.data.default_root_pose.torch.clone() root_pose[:, :3] += origins robot.write_root_pose_to_sim_index(root_pose=root_pose) - root_vel = wp.to_torch(robot.data.default_root_vel).clone() + root_vel = robot.data.default_root_vel.torch.clone() robot.write_root_velocity_to_sim_index(root_velocity=root_vel) # set joint positions with some noise joint_pos, joint_vel = ( - wp.to_torch(robot.data.default_joint_pos).clone(), - wp.to_torch(robot.data.default_joint_vel).clone(), + robot.data.default_joint_pos.torch.clone(), + robot.data.default_joint_vel.torch.clone(), ) joint_pos += torch.rand_like(joint_pos) * 0.1 robot.write_joint_position_to_sim_index(position=joint_pos) @@ -109,7 +108,7 @@ def run_simulator(sim: sim_utils.SimulationContext, entities: dict[str, Articula print("[INFO]: Resetting robot state...") # Apply random action # -- generate random joint efforts - efforts = torch.randn_like(wp.to_torch(robot.data.joint_pos)) * 5.0 + efforts = torch.randn_like(robot.data.joint_pos.torch) * 5.0 # -- apply action to the robot robot.set_joint_effort_target_index(target=efforts) # -- write data to sim diff --git a/scripts/tutorials/01_assets/run_deformable_object.py b/scripts/tutorials/01_assets/run_deformable_object.py index 4d2583974b96..368fe7c3358e 100644 --- a/scripts/tutorials/01_assets/run_deformable_object.py +++ b/scripts/tutorials/01_assets/run_deformable_object.py @@ -36,7 +36,6 @@ """Rest everything follows.""" import torch -import warp as wp from isaaclab_physx.assets import DeformableObject, DeformableObjectCfg # deformables supported in PhysX @@ -98,7 +97,7 @@ def run_simulator(sim: sim_utils.SimulationContext, entities: dict, origins: tor count = 0 # Nodal kinematic targets of the deformable bodies - nodal_kinematic_target = wp.to_torch(cube_object.data.nodal_kinematic_target).clone() + nodal_kinematic_target = cube_object.data.nodal_kinematic_target.torch.clone() # Simulate physics while simulation_app.is_running(): @@ -108,7 +107,7 @@ def run_simulator(sim: sim_utils.SimulationContext, entities: dict, origins: tor count = 0 # reset the nodal state of the object - nodal_state = wp.to_torch(cube_object.data.default_nodal_state_w).clone() + nodal_state = cube_object.data.default_nodal_state_w.torch.clone() # apply random pose to the object pos_w = torch.rand(cube_object.num_instances, 3, device=sim.device) * 0.1 + origins quat_w = math_utils.random_orientation(cube_object.num_instances, device=sim.device) @@ -150,9 +149,7 @@ def run_simulator(sim: sim_utils.SimulationContext, entities: dict, origins: tor # print the root positions every second if count % int(1.0 / sim_dt) == 0: - print( - f"Time {sim_time:.2f}s: \tRoot position (in world): {wp.to_torch(cube_object.data.root_pos_w)[:, :3]}" - ) + print(f"Time {sim_time:.2f}s: \tRoot position (in world): {cube_object.data.root_pos_w.torch[:, :3]}") def main(): diff --git a/scripts/tutorials/01_assets/run_rigid_object.py b/scripts/tutorials/01_assets/run_rigid_object.py index 6cce62d774bf..232e24f9743f 100644 --- a/scripts/tutorials/01_assets/run_rigid_object.py +++ b/scripts/tutorials/01_assets/run_rigid_object.py @@ -34,7 +34,6 @@ """Rest everything follows.""" import torch -import warp as wp import isaaclab.sim as sim_utils import isaaclab.utils.math as math_utils @@ -95,7 +94,7 @@ def run_simulator(sim: sim_utils.SimulationContext, entities: dict[str, RigidObj sim_time = 0.0 count = 0 # reset root state - root_pose = wp.to_torch(cone_object.data.default_root_pose).clone() + root_pose = cone_object.data.default_root_pose.torch.clone() # sample a random position on a cylinder around the origins root_pose[:, :3] += origins root_pose[:, :3] += math_utils.sample_cylinder( @@ -103,7 +102,7 @@ def run_simulator(sim: sim_utils.SimulationContext, entities: dict[str, RigidObj ) # write root state to simulation cone_object.write_root_pose_to_sim_index(root_pose=root_pose) - root_vel = wp.to_torch(cone_object.data.default_root_vel).clone() + root_vel = cone_object.data.default_root_vel.torch.clone() cone_object.write_root_velocity_to_sim_index(root_velocity=root_vel) # reset buffers cone_object.reset() @@ -120,7 +119,7 @@ def run_simulator(sim: sim_utils.SimulationContext, entities: dict[str, RigidObj cone_object.update(sim_dt) # print the root position if count % 50 == 0: - print(f"Root position (in world): {wp.to_torch(cone_object.data.root_pos_w)}") + print(f"Root position (in world): {cone_object.data.root_pos_w.torch}") def main(): diff --git a/scripts/tutorials/01_assets/run_surface_gripper.py b/scripts/tutorials/01_assets/run_surface_gripper.py index cb334f851e4c..5488232a5d42 100644 --- a/scripts/tutorials/01_assets/run_surface_gripper.py +++ b/scripts/tutorials/01_assets/run_surface_gripper.py @@ -108,15 +108,15 @@ def run_simulator( # root state # we offset the root state by the origin since the states are written in simulation world frame # if this is not done, then the robots will be spawned at the (0, 0, 0) of the simulation world - root_pose = wp.to_torch(robot.data.default_root_pose).clone() + root_pose = robot.data.default_root_pose.torch.clone() root_pose[:, :3] += origins robot.write_root_pose_to_sim_index(root_pose=root_pose) - root_vel = wp.to_torch(robot.data.default_root_vel).clone() + root_vel = robot.data.default_root_vel.torch.clone() robot.write_root_velocity_to_sim_index(root_velocity=root_vel) # set joint positions with some noise joint_pos, joint_vel = ( - wp.to_torch(robot.data.default_joint_pos).clone(), - wp.to_torch(robot.data.default_joint_vel).clone(), + robot.data.default_joint_pos.torch.clone(), + robot.data.default_joint_vel.torch.clone(), ) joint_pos += torch.rand_like(joint_pos) * 0.1 robot.write_joint_position_to_sim_index(position=joint_pos) diff --git a/scripts/tutorials/02_scene/create_scene.py b/scripts/tutorials/02_scene/create_scene.py index 7c578e4366b3..1d4578ca256c 100644 --- a/scripts/tutorials/02_scene/create_scene.py +++ b/scripts/tutorials/02_scene/create_scene.py @@ -34,7 +34,6 @@ """Rest everything follows.""" import torch -import warp as wp import isaaclab.sim as sim_utils from isaaclab.assets import ArticulationCfg, AssetBaseCfg @@ -82,15 +81,15 @@ def run_simulator(sim: sim_utils.SimulationContext, scene: InteractiveScene): # root state # we offset the root state by the origin since the states are written in simulation world frame # if this is not done, then the robots will be spawned at the (0, 0, 0) of the simulation world - root_pose = wp.to_torch(robot.data.default_root_pose).clone() + root_pose = robot.data.default_root_pose.torch.clone() root_pose[:, :3] += scene.env_origins robot.write_root_pose_to_sim_index(root_pose=root_pose) - root_vel = wp.to_torch(robot.data.default_root_vel).clone() + root_vel = robot.data.default_root_vel.torch.clone() robot.write_root_velocity_to_sim_index(root_velocity=root_vel) # set joint positions with some noise joint_pos, joint_vel = ( - wp.to_torch(robot.data.default_joint_pos).clone(), - wp.to_torch(robot.data.default_joint_vel).clone(), + robot.data.default_joint_pos.torch.clone(), + robot.data.default_joint_vel.torch.clone(), ) joint_pos += torch.rand_like(joint_pos) * 0.1 robot.write_joint_position_to_sim_index(position=joint_pos) @@ -100,7 +99,7 @@ def run_simulator(sim: sim_utils.SimulationContext, scene: InteractiveScene): print("[INFO]: Resetting robot state...") # Apply random action # -- generate random joint efforts - efforts = torch.randn_like(wp.to_torch(robot.data.joint_pos)) * 5.0 + efforts = torch.randn_like(robot.data.joint_pos.torch) * 5.0 # -- apply action to the robot robot.set_joint_effort_target_index(target=efforts) # -- write data to sim diff --git a/scripts/tutorials/03_envs/create_cube_base_env.py b/scripts/tutorials/03_envs/create_cube_base_env.py index c5850905bd3b..f6629c6bbb82 100644 --- a/scripts/tutorials/03_envs/create_cube_base_env.py +++ b/scripts/tutorials/03_envs/create_cube_base_env.py @@ -51,7 +51,6 @@ """Rest everything follows.""" import torch -import warp as wp import isaaclab.envs.mdp as mdp import isaaclab.sim as sim_utils @@ -129,8 +128,8 @@ def process_actions(self, actions: torch.Tensor): def apply_actions(self): # implement a PD controller to track the target position - pos_error = self._processed_actions - (wp.to_torch(self._asset.data.root_pos_w) - self._env.scene.env_origins) - vel_error = -wp.to_torch(self._asset.data.root_lin_vel_w) + pos_error = self._processed_actions - (self._asset.data.root_pos_w.torch - self._env.scene.env_origins) + vel_error = -self._asset.data.root_lin_vel_w.torch # set velocity targets self._vel_command[:, :3] = self.p_gain * pos_error + self.d_gain * vel_error self._asset.write_root_velocity_to_sim_index(root_velocity=self._vel_command) @@ -158,7 +157,7 @@ def base_position(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg) -> torch.Tens """Root linear velocity in the asset's root frame.""" # extract the used quantities (to enable type-hinting) asset: RigidObject = env.scene[asset_cfg.name] - return wp.to_torch(asset.data.root_pos_w) - env.scene.env_origins + return asset.data.root_pos_w.torch - env.scene.env_origins ## diff --git a/scripts/tutorials/04_sensors/add_sensors_on_robot.py b/scripts/tutorials/04_sensors/add_sensors_on_robot.py index f5e3a19c0bec..4c143318f27a 100644 --- a/scripts/tutorials/04_sensors/add_sensors_on_robot.py +++ b/scripts/tutorials/04_sensors/add_sensors_on_robot.py @@ -40,7 +40,6 @@ """Rest everything follows.""" import torch -import warp as wp import isaaclab.sim as sim_utils from isaaclab.assets import ArticulationCfg, AssetBaseCfg @@ -112,15 +111,15 @@ def run_simulator(sim: sim_utils.SimulationContext, scene: InteractiveScene): # root state # we offset the root state by the origin since the states are written in simulation world frame # if this is not done, then the robots will be spawned at the (0, 0, 0) of the simulation world - root_pose = wp.to_torch(scene["robot"].data.default_root_pose).clone() + root_pose = scene["robot"].data.default_root_pose.torch.clone() root_pose[:, :3] += scene.env_origins scene["robot"].write_root_pose_to_sim_index(root_pose=root_pose) - root_vel = wp.to_torch(scene["robot"].data.default_root_vel).clone() + root_vel = scene["robot"].data.default_root_vel.torch.clone() scene["robot"].write_root_velocity_to_sim_index(root_velocity=root_vel) # set joint positions with some noise joint_pos, joint_vel = ( - wp.to_torch(scene["robot"].data.default_joint_pos).clone(), - wp.to_torch(scene["robot"].data.default_joint_vel).clone(), + scene["robot"].data.default_joint_pos.torch.clone(), + scene["robot"].data.default_joint_vel.torch.clone(), ) joint_pos += torch.rand_like(joint_pos) * 0.1 scene["robot"].write_joint_position_to_sim_index(position=joint_pos) @@ -130,7 +129,7 @@ def run_simulator(sim: sim_utils.SimulationContext, scene: InteractiveScene): print("[INFO]: Resetting robot state...") # Apply default actions to the robot # -- generate actions/commands - targets = wp.to_torch(scene["robot"].data.default_joint_pos) + targets = scene["robot"].data.default_joint_pos.torch # -- apply action to the robot scene["robot"].set_joint_position_target_index(target=targets) # -- write data to sim @@ -152,7 +151,7 @@ def run_simulator(sim: sim_utils.SimulationContext, scene: InteractiveScene): print(scene["height_scanner"]) print( "Received max height value: ", - torch.max(wp.to_torch(scene["height_scanner"].data.ray_hits_w)[..., -1]).item(), + torch.max(scene["height_scanner"].data.ray_hits_w.torch[..., -1]).item(), ) print("-------------------------------") print(scene["contact_forces"]) diff --git a/scripts/tutorials/04_sensors/run_frame_transformer.py b/scripts/tutorials/04_sensors/run_frame_transformer.py index 9370179cb213..980ced307621 100644 --- a/scripts/tutorials/04_sensors/run_frame_transformer.py +++ b/scripts/tutorials/04_sensors/run_frame_transformer.py @@ -35,7 +35,6 @@ import math import torch -import warp as wp import isaacsim.util.debug_draw._debug_draw as omni_debug_draw @@ -124,7 +123,7 @@ def run_simulator(sim: sim_utils.SimulationContext, scene_entities: dict): # Simulate physics while simulation_app.is_running(): # perform this loop at policy control freq (50 Hz) - robot.set_joint_position_target_index(target=wp.to_torch(robot.data.default_joint_pos).clone()) + robot.set_joint_position_target_index(target=robot.data.default_joint_pos.torch.clone()) robot.write_data_to_sim() # perform step sim.step() @@ -147,10 +146,10 @@ def run_simulator(sim: sim_utils.SimulationContext, scene_entities: dict): print(f"Displaying Frame ID {frame_index}: {frame_names[frame_index]}") # visualize frame - source_pos = wp.to_torch(frame_transformer.data.source_pos_w) - source_quat = wp.to_torch(frame_transformer.data.source_quat_w) - target_pos = wp.to_torch(frame_transformer.data.target_pos_w)[:, frame_index] - target_quat = wp.to_torch(frame_transformer.data.target_quat_w)[:, frame_index] + source_pos = frame_transformer.data.source_pos_w.torch + source_quat = frame_transformer.data.source_quat_w.torch + target_pos = frame_transformer.data.target_pos_w.torch[:, frame_index] + target_quat = frame_transformer.data.target_quat_w.torch[:, frame_index] # draw the frames transform_visualizer.visualize( torch.cat([source_pos, target_pos], dim=0), torch.cat([source_quat, target_quat], dim=0) diff --git a/scripts/tutorials/04_sensors/run_ray_caster.py b/scripts/tutorials/04_sensors/run_ray_caster.py index ff66ff9a0fd2..ea6f946d8c5f 100644 --- a/scripts/tutorials/04_sensors/run_ray_caster.py +++ b/scripts/tutorials/04_sensors/run_ray_caster.py @@ -32,7 +32,6 @@ """Rest everything follows.""" import torch -import warp as wp import isaaclab.sim as sim_utils from isaaclab.assets import RigidObject, RigidObjectCfg @@ -98,9 +97,9 @@ def run_simulator(sim: sim_utils.SimulationContext, scene_entities: dict): balls: RigidObject = scene_entities["balls"] # define an initial position of the sensor - ball_default_pose = wp.to_torch(balls.data.default_root_pose).clone() + ball_default_pose = balls.data.default_root_pose.torch.clone() ball_default_pose[:, :3] = torch.rand_like(ball_default_pose[:, :3]) * 10 - ball_default_vel = wp.to_torch(balls.data.default_root_vel).clone() + ball_default_vel = balls.data.default_root_vel.torch.clone() # Create a counter for resetting the scene step_count = 0 @@ -120,7 +119,7 @@ def run_simulator(sim: sim_utils.SimulationContext, scene_entities: dict): # Update the ray-caster with Timer( f"Ray-caster update with {4} x {ray_caster.num_rays} rays with max height of" - f" {torch.max(wp.to_torch(ray_caster.data.pos_w)).item():.2f}" + f" {torch.max(ray_caster.data.pos_w.torch).item():.2f}" ): ray_caster.update(dt=sim.get_physics_dt(), force_recompute=True) # Update counter diff --git a/scripts/tutorials/05_controllers/run_diff_ik.py b/scripts/tutorials/05_controllers/run_diff_ik.py index f4a5c7dcf603..2922265e7b9b 100644 --- a/scripts/tutorials/05_controllers/run_diff_ik.py +++ b/scripts/tutorials/05_controllers/run_diff_ik.py @@ -20,8 +20,6 @@ import argparse -import warp as wp - from isaaclab.app import AppLauncher # add argparse arguments @@ -147,8 +145,8 @@ def run_simulator(sim: sim_utils.SimulationContext, scene: InteractiveScene): # reset time count = 0 # reset joint state - joint_pos = wp.to_torch(robot.data.default_joint_pos).clone() - joint_vel = wp.to_torch(robot.data.default_joint_vel).clone() + joint_pos = robot.data.default_joint_pos.torch.clone() + joint_vel = robot.data.default_joint_vel.torch.clone() robot.write_joint_position_to_sim_index(position=joint_pos) robot.write_joint_velocity_to_sim_index(velocity=joint_vel) robot.reset() @@ -163,9 +161,9 @@ def run_simulator(sim: sim_utils.SimulationContext, scene: InteractiveScene): else: # obtain quantities from simulation jacobian = robot.root_view.get_jacobians()[:, ee_jacobi_idx, :, robot_entity_cfg.joint_ids] - ee_pose_w = wp.to_torch(robot.data.body_pose_w)[:, robot_entity_cfg.body_ids[0]] - root_pose_w = wp.to_torch(robot.data.root_pose_w) - joint_pos = wp.to_torch(robot.data.joint_pos)[:, robot_entity_cfg.joint_ids] + ee_pose_w = robot.data.body_pose_w.torch[:, robot_entity_cfg.body_ids[0]] + root_pose_w = robot.data.root_pose_w.torch + joint_pos = robot.data.joint_pos.torch[:, robot_entity_cfg.joint_ids] # compute frame in root frame ee_pos_b, ee_quat_b = subtract_frame_transforms( root_pose_w[:, 0:3], root_pose_w[:, 3:7], ee_pose_w[:, 0:3], ee_pose_w[:, 3:7] @@ -184,7 +182,7 @@ def run_simulator(sim: sim_utils.SimulationContext, scene: InteractiveScene): scene.update(sim_dt) # obtain quantities from simulation - ee_pose_w = wp.to_torch(robot.data.body_state_w)[:, robot_entity_cfg.body_ids[0], 0:7] + ee_pose_w = robot.data.body_state_w.torch[:, robot_entity_cfg.body_ids[0], 0:7] # update marker positions ee_marker.visualize(ee_pose_w[:, 0:3], ee_pose_w[:, 3:7]) goal_marker.visualize(ik_commands[:, 0:3] + scene.env_origins, ik_commands[:, 3:7]) diff --git a/scripts/tutorials/05_controllers/run_osc.py b/scripts/tutorials/05_controllers/run_osc.py index 841ee0c9dd68..09e16c40c70e 100644 --- a/scripts/tutorials/05_controllers/run_osc.py +++ b/scripts/tutorials/05_controllers/run_osc.py @@ -20,8 +20,6 @@ import argparse -import warp as wp - from isaaclab.app import AppLauncher # add argparse arguments @@ -181,7 +179,7 @@ def run_simulator(sim: sim_utils.SimulationContext, scene: InteractiveScene): robot.update(dt=sim_dt) # Get the center of the robot soft joint limits - joint_centers = torch.mean(wp.to_torch(robot.data.soft_joint_pos_limits)[:, arm_joint_ids, :], dim=-1) + joint_centers = torch.mean(robot.data.soft_joint_pos_limits.torch[:, arm_joint_ids, :], dim=-1) # get the updated states ( @@ -215,8 +213,8 @@ def run_simulator(sim: sim_utils.SimulationContext, scene: InteractiveScene): # reset every 500 steps if count % 500 == 0: # reset joint state to default - default_joint_pos = wp.to_torch(robot.data.default_joint_pos).clone() - default_joint_vel = wp.to_torch(robot.data.default_joint_vel).clone() + default_joint_pos = robot.data.default_joint_pos.torch.clone() + default_joint_vel = robot.data.default_joint_vel.torch.clone() robot.write_joint_position_to_sim_index(position=default_joint_pos) robot.write_joint_velocity_to_sim_index(velocity=default_joint_vel) robot.set_joint_effort_target_index(target=zero_joint_efforts) # Set zero torques in the initial step @@ -321,30 +319,26 @@ def update_states( gravity = robot.root_view.get_gravity_compensation_forces()[:, arm_joint_ids] # Convert the Jacobian from world to root frame jacobian_b = jacobian_w.clone() - root_rot_matrix = matrix_from_quat(quat_inv(wp.to_torch(robot.data.root_quat_w))) + root_rot_matrix = matrix_from_quat(quat_inv(robot.data.root_quat_w.torch)) jacobian_b[:, :3, :] = torch.bmm(root_rot_matrix, jacobian_b[:, :3, :]) jacobian_b[:, 3:, :] = torch.bmm(root_rot_matrix, jacobian_b[:, 3:, :]) # Compute current pose of the end-effector - root_pos_w = wp.to_torch(robot.data.root_pos_w) - root_quat_w = wp.to_torch(robot.data.root_quat_w) - ee_pos_w = wp.to_torch(robot.data.body_pos_w)[:, ee_frame_idx] - ee_quat_w = wp.to_torch(robot.data.body_quat_w)[:, ee_frame_idx] + root_pos_w = robot.data.root_pos_w.torch + root_quat_w = robot.data.root_quat_w.torch + ee_pos_w = robot.data.body_pos_w.torch[:, ee_frame_idx] + ee_quat_w = robot.data.body_quat_w.torch[:, ee_frame_idx] ee_pos_b, ee_quat_b = subtract_frame_transforms(root_pos_w, root_quat_w, ee_pos_w, ee_quat_w) root_pose_w = torch.cat([root_pos_w, root_quat_w], dim=-1) ee_pose_w = torch.cat([ee_pos_w, ee_quat_w], dim=-1) ee_pose_b = torch.cat([ee_pos_b, ee_quat_b], dim=-1) # Compute the current velocity of the end-effector - ee_vel_w = wp.to_torch(robot.data.body_vel_w)[ - :, ee_frame_idx, : - ] # Extract end-effector velocity in the world frame - root_vel_w = wp.to_torch(robot.data.root_vel_w) # Extract root velocity in the world frame + ee_vel_w = robot.data.body_vel_w.torch[:, ee_frame_idx, :] # Extract end-effector velocity in the world frame + root_vel_w = robot.data.root_vel_w.torch # Extract root velocity in the world frame relative_vel_w = ee_vel_w - root_vel_w # Compute the relative velocity in the world frame - ee_lin_vel_b = quat_apply_inverse( - wp.to_torch(robot.data.root_quat_w), relative_vel_w[:, 0:3] - ) # From world to root frame - ee_ang_vel_b = quat_apply_inverse(wp.to_torch(robot.data.root_quat_w), relative_vel_w[:, 3:6]) + ee_lin_vel_b = quat_apply_inverse(robot.data.root_quat_w.torch, relative_vel_w[:, 0:3]) # From world to root frame + ee_ang_vel_b = quat_apply_inverse(robot.data.root_quat_w.torch, relative_vel_w[:, 3:6]) ee_vel_b = torch.cat([ee_lin_vel_b, ee_ang_vel_b], dim=-1) # Calculate the contact force @@ -359,8 +353,8 @@ def update_states( ee_force_b = ee_force_w # Get joint positions and velocities - joint_pos = wp.to_torch(robot.data.joint_pos)[:, arm_joint_ids] - joint_vel = wp.to_torch(robot.data.joint_vel)[:, arm_joint_ids] + joint_pos = robot.data.joint_pos.torch[:, arm_joint_ids] + joint_vel = robot.data.joint_vel.torch[:, arm_joint_ids] return ( jacobian_b, diff --git a/source/isaaclab/config/extension.toml b/source/isaaclab/config/extension.toml index 0f93bbd09d09..54d18e0a4e93 100644 --- a/source/isaaclab/config/extension.toml +++ b/source/isaaclab/config/extension.toml @@ -1,7 +1,7 @@ [package] # Note: Semantic Versioning is used: https://semver.org/ -version = "4.6.16" +version = "4.6.21" # Description title = "Isaac Lab framework for Robot Learning" diff --git a/source/isaaclab/docs/CHANGELOG.rst b/source/isaaclab/docs/CHANGELOG.rst index fdfcc95e4738..ddb86e5d0105 100644 --- a/source/isaaclab/docs/CHANGELOG.rst +++ b/source/isaaclab/docs/CHANGELOG.rst @@ -1,7 +1,18 @@ Changelog --------- -4.6.16 (2026-04-24) +4.6.21 (2026-04-27) +~~~~~~~~~~~~~~~~~~~ + +Fixed +^^^^^ + +* Fixed ProxyArray migration guidance and static regression checks to cover + deprecated ``wp.to_torch(proxy_array)`` usage and direct tensor/wp.array + method calls on data properties. + + +4.6.20 (2026-04-27) ~~~~~~~~~~~~~~~~~~~ Added @@ -12,9 +23,17 @@ Added (typically an upstream library config object). Raises :class:`AttributeError` if the target is missing a declared field, so upstream renames surface at startup instead of as silent no-ops. +* Added an opt-in runtime detector on :class:`~isaaclab.utils.warp.ProxyArray` + that emits a :class:`UserWarning` with the call site (``stacklevel=2``) on + every ``.torch`` read of a ``wp.quatf``-typed array when the + ``WARN_ON_TORCH_QUATF_ACCESS`` environment variable is set to ``"1"``. + Helps users find code that still assumes Isaac Lab 2.x's ``(w, x, y, z)`` + quaternion convention after the migration to Isaac Lab 3.x's + ``(x, y, z, w)`` convention. Documented in the Isaac Lab 3.0 migration + guide as a complement to the source-level quaternion finder tool. -4.6.15 (2026-04-24) +4.6.19 (2026-04-27) ~~~~~~~~~~~~~~~~~~~ Changed @@ -25,7 +44,7 @@ Changed ``@optional_method``. -4.6.14 (2026-04-24) +4.6.18 (2026-04-27) ~~~~~~~~~~~~~~~~~~~ Changed @@ -55,7 +74,7 @@ Fixed to Replicator so the resolver returns a valid path. -4.6.13 (2026-04-22) +4.6.17 (2026-04-27) ~~~~~~~~~~~~~~~~~~~ Changed @@ -89,6 +108,129 @@ Removed for similar debugging workflows. +4.6.16 (2026-04-24) +~~~~~~~~~~~~~~~~~~~ + +Changed +^^^^^^^ + +* Renamed the ``TorchArray`` wrapper to + :class:`~isaaclab.utils.warp.ProxyArray` and its module from + ``isaaclab.utils.warp.torch_array`` to + ``isaaclab.utils.warp.proxy_array``. The new name better describes + the class as a proxy around a :class:`wp.array` exposing both + ``.warp`` and ``.torch`` accessors, rather than a torch-first type. + No backward-compatibility alias — downstream code must update + imports and type hints. + +Removed +^^^^^^^ + +* Removed :meth:`~isaaclab.utils.warp.ProxyArray.rebind`. ``ProxyArray`` + instances are now immutable after construction: assigning to any + attribute other than the internal ``_torch_cache`` raises + :class:`AttributeError`. Call sites that previously rebound a + sim-bound wrapper (Newton / PhysX asset data classes when the solver + re-creates its arrays) now assign a fresh + ``ProxyArray(new_buf)`` instead. The old pattern was unsafe — + callers holding a stale ``.torch`` reference read the old memory + silently after ``rebind``. Forcing a new wrapper per rebind makes + the stale-reference hazard local and explicit. + + .. code-block:: python + + # Before (4.6.15) + self._joint_pos_ta.rebind(self._sim_bind_joint_pos) + # After (4.6.16) + self._joint_pos_ta = ProxyArray(self._sim_bind_joint_pos) + + +4.6.15 (2026-04-24) +~~~~~~~~~~~~~~~~~~~ + +Changed +^^^^^^^ + +* :meth:`~isaaclab.sim.views.BaseFrameView.get_world_poses` and + :meth:`~isaaclab.sim.views.BaseFrameView.get_local_poses` now return + a pair of :class:`~isaaclab.utils.warp.ProxyArray` wrappers instead + of raw ``wp.array``. :class:`~isaaclab.sim.views.UsdFrameView`, + :class:`~isaaclab_physx.sim.views.FabricFrameView`, and + :class:`~isaaclab_newton.sim.views.NewtonSiteFrameView` were updated + accordingly. Calls without ``indices`` reuse a cached ProxyArray + wrapping the backend's stable buffer (Fabric/Newton); calls with + ``indices`` and all calls on USD construct a fresh ProxyArray per + invocation. + + **Breaking change** — call sites that wrapped the return in + ``wp.to_torch(...)`` or passed the result directly to a kernel must + migrate: + + .. code-block:: python + + # Before (4.6.14) + pos_wp, quat_wp = view.get_world_poses(indices) + tensor = wp.to_torch(pos_wp) + # After (4.6.15) + pos_w, quat_w = view.get_world_poses(indices) + tensor = pos_w.torch + # ... and when passing to a warp kernel: + wp.launch(..., inputs=[pos_w.warp, ...]) + + +4.6.14 (2026-04-24) +~~~~~~~~~~~~~~~~~~~ + +Changed +^^^^^^^ + +* :class:`~isaaclab.sensors.RayCasterData` and + :class:`~isaaclab.sensors.MultiMeshRayCasterData` properties + :attr:`~isaaclab.sensors.RayCasterData.pos_w`, + :attr:`~isaaclab.sensors.RayCasterData.quat_w`, and + :attr:`~isaaclab.sensors.RayCasterData.ray_hits_w` now return + :class:`~isaaclab.utils.warp.ProxyArray` instead of raw ``wp.array``. + Use ``.torch`` for a cached zero-copy ``torch.Tensor`` view or ``.warp`` + for the underlying ``wp.array``. + + **Breaking change** — call sites that previously wrapped these in + ``wp.to_torch(...)`` must migrate to the ``.torch`` accessor: + + .. code-block:: python + + # Before (4.6.13) + hits = wp.to_torch(sensor.data.ray_hits_w) + # After (4.6.14) + hits = sensor.data.ray_hits_w.torch + + +4.6.13 (2026-04-23) +~~~~~~~~~~~~~~~~~~~ + +Added +^^^^^ + +* Added :class:`~isaaclab.utils.warp.ProxyArray`, a warp-first dual-access array + that provides explicit ``.torch`` and ``.warp`` accessors for seamless + interoperability between warp and PyTorch workflows. + +Changed +^^^^^^^ + +* All properties on the following base data classes now return + :class:`~isaaclab.utils.warp.ProxyArray` instead of raw ``wp.array``: + :class:`~isaaclab.assets.articulation.BaseArticulationData`, + :class:`~isaaclab.assets.rigid_object.BaseRigidObjectData`, + :class:`~isaaclab.assets.rigid_object_collection.BaseRigidObjectCollectionData`, + :class:`~isaaclab.sensors.contact_sensor.BaseContactSensorData`, + :class:`~isaaclab.sensors.frame_transformer.BaseFrameTransformerData`, + :class:`~isaaclab.sensors.imu.BaseImuData`, and + :class:`~isaaclab.sensors.pva.BasePvaData`. + Use ``.torch`` for a cached zero-copy ``torch.Tensor`` view, or ``.warp`` for + the underlying ``wp.array``. Implicit torch operations (arithmetic, + ``torch.*`` functions) work during the deprecation period but emit a warning. + + 4.6.12 (2026-04-23) ~~~~~~~~~~~~~~~~~~~ diff --git a/source/isaaclab/isaaclab/assets/articulation/base_articulation_data.py b/source/isaaclab/isaaclab/assets/articulation/base_articulation_data.py index 01dab490183a..fdd79ce6d474 100644 --- a/source/isaaclab/isaaclab/assets/articulation/base_articulation_data.py +++ b/source/isaaclab/isaaclab/assets/articulation/base_articulation_data.py @@ -8,6 +8,8 @@ import warp as wp +from isaaclab.utils.warp import ProxyArray + class BaseArticulationData(ABC): """Data container for an articulation. @@ -63,7 +65,7 @@ def update(self, dt: float) -> None: @property @abstractmethod - def default_root_pose(self) -> wp.array: + def default_root_pose(self) -> ProxyArray: """Default root pose ``[pos, quat]`` in the local environment frame. The position and quaternion are of the articulation root's actor frame. Shape is (num_instances), @@ -73,7 +75,7 @@ def default_root_pose(self) -> wp.array: @property @abstractmethod - def default_root_vel(self) -> wp.array: + def default_root_vel(self) -> ProxyArray: """Default root velocity ``[lin_vel, ang_vel]`` in the local environment frame. The linear and angular velocities are of the articulation root's center of mass frame. @@ -83,13 +85,13 @@ def default_root_vel(self) -> wp.array: @property @abstractmethod - def default_root_state(self) -> wp.array: + def default_root_state(self) -> ProxyArray: """Deprecated, same as :attr:`default_root_pose` and :attr:`default_root_vel`.""" raise NotImplementedError @property @abstractmethod - def default_joint_pos(self) -> wp.array: + def default_joint_pos(self) -> ProxyArray: """Default joint positions of all joints. Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to (num_instances, num_joints). @@ -100,7 +102,7 @@ def default_joint_pos(self) -> wp.array: @property @abstractmethod - def default_joint_vel(self) -> wp.array: + def default_joint_vel(self) -> ProxyArray: """Default joint velocities of all joints. Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to (num_instances, num_joints). @@ -115,7 +117,7 @@ def default_joint_vel(self) -> wp.array: @property @abstractmethod - def joint_pos_target(self) -> wp.array: + def joint_pos_target(self) -> ProxyArray: """Joint position targets commanded by the user. Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to (num_instances, num_joints). @@ -128,7 +130,7 @@ def joint_pos_target(self) -> wp.array: @property @abstractmethod - def joint_vel_target(self) -> wp.array: + def joint_vel_target(self) -> ProxyArray: """Joint velocity targets commanded by the user. Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to (num_instances, num_joints). @@ -141,7 +143,7 @@ def joint_vel_target(self) -> wp.array: @property @abstractmethod - def joint_effort_target(self) -> wp.array: + def joint_effort_target(self) -> ProxyArray: """Joint effort targets commanded by the user. Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to (num_instances, num_joints). @@ -158,7 +160,7 @@ def joint_effort_target(self) -> wp.array: @property @abstractmethod - def computed_torque(self) -> wp.array: + def computed_torque(self) -> ProxyArray: """Joint torques computed from the actuator model (before clipping). Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to (num_instances, num_joints). @@ -171,7 +173,7 @@ def computed_torque(self) -> wp.array: @property @abstractmethod - def applied_torque(self) -> wp.array: + def applied_torque(self) -> ProxyArray: """Joint torques applied from the actuator model (after clipping). Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to (num_instances, num_joints). @@ -187,7 +189,7 @@ def applied_torque(self) -> wp.array: @property @abstractmethod - def joint_stiffness(self) -> wp.array: + def joint_stiffness(self) -> ProxyArray: """Joint stiffness provided to the simulation. Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to (num_instances, num_joints). @@ -198,7 +200,7 @@ def joint_stiffness(self) -> wp.array: @property @abstractmethod - def joint_damping(self) -> wp.array: + def joint_damping(self) -> ProxyArray: """Joint damping provided to the simulation. Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to (num_instances, num_joints). @@ -209,7 +211,7 @@ def joint_damping(self) -> wp.array: @property @abstractmethod - def joint_armature(self) -> wp.array: + def joint_armature(self) -> ProxyArray: """Joint armature provided to the simulation. Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to (num_instances, num_joints). @@ -218,7 +220,7 @@ def joint_armature(self) -> wp.array: @property @abstractmethod - def joint_friction_coeff(self) -> wp.array: + def joint_friction_coeff(self) -> ProxyArray: """Joint static friction coefficient provided to the simulation. Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to (num_instances, num_joints). @@ -227,7 +229,7 @@ def joint_friction_coeff(self) -> wp.array: @property @abstractmethod - def joint_pos_limits(self) -> wp.array: + def joint_pos_limits(self) -> ProxyArray: """Joint position limits provided to the simulation. Shape is (num_instances, num_joints, 2), dtype = wp.vec2f. In torch this resolves to @@ -239,7 +241,7 @@ def joint_pos_limits(self) -> wp.array: @property @abstractmethod - def joint_vel_limits(self) -> wp.array: + def joint_vel_limits(self) -> ProxyArray: """Joint maximum velocity provided to the simulation. Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to (num_instances, num_joints). @@ -248,7 +250,7 @@ def joint_vel_limits(self) -> wp.array: @property @abstractmethod - def joint_effort_limits(self) -> wp.array: + def joint_effort_limits(self) -> ProxyArray: """Joint maximum effort provided to the simulation. Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to (num_instances, num_joints). @@ -261,7 +263,7 @@ def joint_effort_limits(self) -> wp.array: @property @abstractmethod - def soft_joint_pos_limits(self) -> wp.array: + def soft_joint_pos_limits(self) -> ProxyArray: r"""Soft joint positions limits for all joints. Shape is (num_instances, num_joints), dtype = wp.vec2f. In torch this resolves to @@ -286,7 +288,7 @@ def soft_joint_pos_limits(self) -> wp.array: @property @abstractmethod - def soft_joint_vel_limits(self) -> wp.array: + def soft_joint_vel_limits(self) -> ProxyArray: """Soft joint velocity limits for all joints. Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to (num_instances, num_joints). @@ -298,7 +300,7 @@ def soft_joint_vel_limits(self) -> wp.array: @property @abstractmethod - def gear_ratio(self) -> wp.array: + def gear_ratio(self) -> ProxyArray: """Gear ratio for relating motor torques to applied Joint torques. Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to (num_instances, num_joints). @@ -311,7 +313,7 @@ def gear_ratio(self) -> wp.array: @property @abstractmethod - def fixed_tendon_stiffness(self) -> wp.array: + def fixed_tendon_stiffness(self) -> ProxyArray: """Fixed tendon stiffness provided to the simulation. Shape is (num_instances, num_fixed_tendons), dtype = wp.float32. In torch this resolves to @@ -321,7 +323,7 @@ def fixed_tendon_stiffness(self) -> wp.array: @property @abstractmethod - def fixed_tendon_damping(self) -> wp.array: + def fixed_tendon_damping(self) -> ProxyArray: """Fixed tendon damping provided to the simulation. Shape is (num_instances, num_fixed_tendons), dtype = wp.float32. In torch this resolves to @@ -331,7 +333,7 @@ def fixed_tendon_damping(self) -> wp.array: @property @abstractmethod - def fixed_tendon_limit_stiffness(self) -> wp.array: + def fixed_tendon_limit_stiffness(self) -> ProxyArray: """Fixed tendon limit stiffness provided to the simulation. Shape is (num_instances, num_fixed_tendons), dtype = wp.float32. In torch this resolves to @@ -341,7 +343,7 @@ def fixed_tendon_limit_stiffness(self) -> wp.array: @property @abstractmethod - def fixed_tendon_rest_length(self) -> wp.array: + def fixed_tendon_rest_length(self) -> ProxyArray: """Fixed tendon rest length provided to the simulation. Shape is (num_instances, num_fixed_tendons), dtype = wp.float32. In torch this resolves to @@ -351,7 +353,7 @@ def fixed_tendon_rest_length(self) -> wp.array: @property @abstractmethod - def fixed_tendon_offset(self) -> wp.array: + def fixed_tendon_offset(self) -> ProxyArray: """Fixed tendon offset provided to the simulation. Shape is (num_instances, num_fixed_tendons), dtype = wp.float32. In torch this resolves to @@ -361,7 +363,7 @@ def fixed_tendon_offset(self) -> wp.array: @property @abstractmethod - def fixed_tendon_pos_limits(self) -> wp.array: + def fixed_tendon_pos_limits(self) -> ProxyArray: """Fixed tendon position limits provided to the simulation. Shape is (num_instances, num_fixed_tendons, 2), dtype = wp.vec2f. In torch this resolves to @@ -375,7 +377,7 @@ def fixed_tendon_pos_limits(self) -> wp.array: @property @abstractmethod - def spatial_tendon_stiffness(self) -> wp.array: + def spatial_tendon_stiffness(self) -> ProxyArray: """Spatial tendon stiffness provided to the simulation. Shape is (num_instances, num_spatial_tendons), dtype = wp.float32. In torch this resolves to @@ -385,7 +387,7 @@ def spatial_tendon_stiffness(self) -> wp.array: @property @abstractmethod - def spatial_tendon_damping(self) -> wp.array: + def spatial_tendon_damping(self) -> ProxyArray: """Spatial tendon damping provided to the simulation. Shape is (num_instances, num_spatial_tendons), dtype = wp.float32. In torch this resolves to @@ -395,7 +397,7 @@ def spatial_tendon_damping(self) -> wp.array: @property @abstractmethod - def spatial_tendon_limit_stiffness(self) -> wp.array: + def spatial_tendon_limit_stiffness(self) -> ProxyArray: """Spatial tendon limit stiffness provided to the simulation. Shape is (num_instances, num_spatial_tendons), dtype = wp.float32. In torch this resolves to @@ -405,7 +407,7 @@ def spatial_tendon_limit_stiffness(self) -> wp.array: @property @abstractmethod - def spatial_tendon_offset(self) -> wp.array: + def spatial_tendon_offset(self) -> ProxyArray: """Spatial tendon offset provided to the simulation. Shape is (num_instances, num_spatial_tendons), dtype = wp.float32. In torch this resolves to @@ -419,7 +421,7 @@ def spatial_tendon_offset(self) -> wp.array: @property @abstractmethod - def root_link_pose_w(self) -> wp.array: + def root_link_pose_w(self) -> ProxyArray: """Root link pose ``[pos, quat]`` in simulation world frame. Shape is (num_instances,), dtype = wp.transformf. In torch this resolves to (num_instances, 7). @@ -431,7 +433,7 @@ def root_link_pose_w(self) -> wp.array: @property @abstractmethod - def root_link_vel_w(self) -> wp.array: + def root_link_vel_w(self) -> ProxyArray: """Root link velocity ``[lin_vel, ang_vel]`` in simulation world frame. Shape is (num_instances,), dtype = wp.spatial_vectorf. In torch this resolves to (num_instances, 6). @@ -443,7 +445,7 @@ def root_link_vel_w(self) -> wp.array: @property @abstractmethod - def root_com_pose_w(self) -> wp.array: + def root_com_pose_w(self) -> ProxyArray: """Root center of mass pose ``[pos, quat]`` in simulation world frame. Shape is (num_instances,), dtype = wp.transformf. In torch this resolves to (num_instances, 7). @@ -455,7 +457,7 @@ def root_com_pose_w(self) -> wp.array: @property @abstractmethod - def root_com_vel_w(self) -> wp.array: + def root_com_vel_w(self) -> ProxyArray: """Root center of mass velocity ``[lin_vel, ang_vel]`` in simulation world frame. Shape is (num_instances,), dtype = wp.spatial_vectorf. In torch this resolves to (num_instances, 6). @@ -467,19 +469,19 @@ def root_com_vel_w(self) -> wp.array: @property @abstractmethod - def root_state_w(self) -> wp.array: + def root_state_w(self) -> ProxyArray: """Deprecated, same as :attr:`root_link_pose_w` and :attr:`root_com_vel_w`.""" raise NotImplementedError @property @abstractmethod - def root_link_state_w(self) -> wp.array: + def root_link_state_w(self) -> ProxyArray: """Deprecated, same as :attr:`root_link_pose_w` and :attr:`root_link_vel_w`.""" raise NotImplementedError @property @abstractmethod - def root_com_state_w(self) -> wp.array: + def root_com_state_w(self) -> ProxyArray: """Deprecated, same as :attr:`root_com_pose_w` and :attr:`root_com_vel_w`.""" raise NotImplementedError @@ -489,7 +491,7 @@ def root_com_state_w(self) -> wp.array: @property @abstractmethod - def body_mass(self) -> wp.array: + def body_mass(self) -> ProxyArray: """Body mass ``wp.float32`` in the world frame. Shape is (num_instances, num_bodies), dtype = wp.float32. In torch this resolves to (num_instances, num_bodies). @@ -498,7 +500,7 @@ def body_mass(self) -> wp.array: @property @abstractmethod - def body_inertia(self) -> wp.array: + def body_inertia(self) -> ProxyArray: """Flattened body inertia in the world frame. Shape is (num_instances, num_bodies, 9), dtype = wp.float32. In torch this resolves to @@ -508,7 +510,7 @@ def body_inertia(self) -> wp.array: @property @abstractmethod - def body_link_pose_w(self) -> wp.array: + def body_link_pose_w(self) -> ProxyArray: """Body link pose ``[pos, quat]`` in simulation world frame. Shape is (num_instances, num_bodies), dtype = wp.transformf. In torch this resolves to @@ -521,7 +523,7 @@ def body_link_pose_w(self) -> wp.array: @property @abstractmethod - def body_link_vel_w(self) -> wp.array: + def body_link_vel_w(self) -> ProxyArray: """Body link velocity ``[lin_vel, ang_vel]`` in simulation world frame. Shape is (num_instances, num_bodies), dtype = wp.spatial_vectorf. In torch this resolves to @@ -534,7 +536,7 @@ def body_link_vel_w(self) -> wp.array: @property @abstractmethod - def body_com_pose_w(self) -> wp.array: + def body_com_pose_w(self) -> ProxyArray: """Body center of mass pose ``[pos, quat]`` in simulation world frame. Shape is (num_instances, num_bodies), dtype = wp.transformf. In torch this resolves to @@ -547,7 +549,7 @@ def body_com_pose_w(self) -> wp.array: @property @abstractmethod - def body_com_vel_w(self) -> wp.array: + def body_com_vel_w(self) -> ProxyArray: """Body center of mass velocity ``[lin_vel, ang_vel]`` in simulation world frame. Shape is (num_instances, num_bodies), dtype = wp.spatial_vectorf. In torch this resolves to @@ -560,25 +562,25 @@ def body_com_vel_w(self) -> wp.array: @property @abstractmethod - def body_state_w(self) -> wp.array: + def body_state_w(self) -> ProxyArray: """Deprecated, same as :attr:`body_link_pose_w` and :attr:`body_com_vel_w`.""" raise NotImplementedError @property @abstractmethod - def body_link_state_w(self) -> wp.array: + def body_link_state_w(self) -> ProxyArray: """Deprecated, same as :attr:`body_link_pose_w` and :attr:`body_link_vel_w`.""" raise NotImplementedError @property @abstractmethod - def body_com_state_w(self) -> wp.array: + def body_com_state_w(self) -> ProxyArray: """Deprecated, same as :attr:`body_com_pose_w` and :attr:`body_com_vel_w`.""" raise NotImplementedError @property @abstractmethod - def body_com_acc_w(self) -> wp.array: + def body_com_acc_w(self) -> ProxyArray: """Acceleration of all bodies center of mass ``[lin_acc, ang_acc]``. Shape is (num_instances, num_bodies), dtype = wp.spatial_vectorf. In torch this resolves to @@ -590,7 +592,7 @@ def body_com_acc_w(self) -> wp.array: @property @abstractmethod - def body_com_pose_b(self) -> wp.array: + def body_com_pose_b(self) -> ProxyArray: """Center of mass pose ``[pos, quat]`` of all bodies in their respective body's link frames. Shape is (num_instances, num_bodies), dtype = wp.transformf. In torch this resolves to @@ -603,7 +605,7 @@ def body_com_pose_b(self) -> wp.array: @property @abstractmethod - def body_incoming_joint_wrench_b(self) -> wp.array: + def body_incoming_joint_wrench_b(self) -> ProxyArray: """Joint reaction wrench applied from body parent to child body in parent body frame. Shape is (num_instances, num_bodies), dtype = wp.spatial_vectorf. In torch this resolves to @@ -624,7 +626,7 @@ def body_incoming_joint_wrench_b(self) -> wp.array: @property @abstractmethod - def joint_pos(self) -> wp.array: + def joint_pos(self) -> ProxyArray: """Joint positions of all joints. Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to @@ -634,7 +636,7 @@ def joint_pos(self) -> wp.array: @property @abstractmethod - def joint_vel(self) -> wp.array: + def joint_vel(self) -> ProxyArray: """Joint velocities of all joints. Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to @@ -644,7 +646,7 @@ def joint_vel(self) -> wp.array: @property @abstractmethod - def joint_acc(self) -> wp.array: + def joint_acc(self) -> ProxyArray: """Joint acceleration of all joints. Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to @@ -658,7 +660,7 @@ def joint_acc(self) -> wp.array: @property @abstractmethod - def projected_gravity_b(self) -> wp.array: + def projected_gravity_b(self) -> ProxyArray: """Projection of the gravity direction on base frame. Shape is (num_instances), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). @@ -667,7 +669,7 @@ def projected_gravity_b(self) -> wp.array: @property @abstractmethod - def heading_w(self) -> wp.array: + def heading_w(self) -> ProxyArray: """Yaw heading of the base frame (in radians). Shape is (num_instances), dtype = wp.float32. In torch this resolves to (num_instances,). @@ -680,7 +682,7 @@ def heading_w(self) -> wp.array: @property @abstractmethod - def root_link_lin_vel_b(self) -> wp.array: + def root_link_lin_vel_b(self) -> ProxyArray: """Root link linear velocity in base frame. Shape is (num_instances), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). @@ -692,7 +694,7 @@ def root_link_lin_vel_b(self) -> wp.array: @property @abstractmethod - def root_link_ang_vel_b(self) -> wp.array: + def root_link_ang_vel_b(self) -> ProxyArray: """Root link angular velocity in base frame. Shape is (num_instances), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). @@ -704,7 +706,7 @@ def root_link_ang_vel_b(self) -> wp.array: @property @abstractmethod - def root_com_lin_vel_b(self) -> wp.array: + def root_com_lin_vel_b(self) -> ProxyArray: """Root center of mass linear velocity in base frame. Shape is (num_instances), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). @@ -716,7 +718,7 @@ def root_com_lin_vel_b(self) -> wp.array: @property @abstractmethod - def root_com_ang_vel_b(self) -> wp.array: + def root_com_ang_vel_b(self) -> ProxyArray: """Root center of mass angular velocity in base frame. Shape is (num_instances), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). @@ -732,7 +734,7 @@ def root_com_ang_vel_b(self) -> wp.array: @property @abstractmethod - def root_link_pos_w(self) -> wp.array: + def root_link_pos_w(self) -> ProxyArray: """Root link position in simulation world frame. Shape is (num_instances), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). @@ -743,7 +745,7 @@ def root_link_pos_w(self) -> wp.array: @property @abstractmethod - def root_link_quat_w(self) -> wp.array: + def root_link_quat_w(self) -> ProxyArray: """Root link orientation (x, y, z, w) in simulation world frame. Shape is (num_instances), dtype = wp.quatf. In torch this resolves to (num_instances, 4). @@ -754,7 +756,7 @@ def root_link_quat_w(self) -> wp.array: @property @abstractmethod - def root_link_lin_vel_w(self) -> wp.array: + def root_link_lin_vel_w(self) -> ProxyArray: """Root linear velocity in simulation world frame. Shape is (num_instances), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). @@ -765,7 +767,7 @@ def root_link_lin_vel_w(self) -> wp.array: @property @abstractmethod - def root_link_ang_vel_w(self) -> wp.array: + def root_link_ang_vel_w(self) -> ProxyArray: """Root link angular velocity in simulation world frame. Shape is (num_instances), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). @@ -776,7 +778,7 @@ def root_link_ang_vel_w(self) -> wp.array: @property @abstractmethod - def root_com_pos_w(self) -> wp.array: + def root_com_pos_w(self) -> ProxyArray: """Root center of mass position in simulation world frame. Shape is (num_instances), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). @@ -787,7 +789,7 @@ def root_com_pos_w(self) -> wp.array: @property @abstractmethod - def root_com_quat_w(self) -> wp.array: + def root_com_quat_w(self) -> ProxyArray: """Root center of mass orientation (x, y, z, w) in simulation world frame. Shape is (num_instances), dtype = wp.quatf. In torch this resolves to (num_instances, 4). @@ -798,7 +800,7 @@ def root_com_quat_w(self) -> wp.array: @property @abstractmethod - def root_com_lin_vel_w(self) -> wp.array: + def root_com_lin_vel_w(self) -> ProxyArray: """Root center of mass linear velocity in simulation world frame. Shape is (num_instances), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). @@ -809,7 +811,7 @@ def root_com_lin_vel_w(self) -> wp.array: @property @abstractmethod - def root_com_ang_vel_w(self) -> wp.array: + def root_com_ang_vel_w(self) -> ProxyArray: """Root center of mass angular velocity in simulation world frame. Shape is (num_instances), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). @@ -820,7 +822,7 @@ def root_com_ang_vel_w(self) -> wp.array: @property @abstractmethod - def body_link_pos_w(self) -> wp.array: + def body_link_pos_w(self) -> ProxyArray: """Positions of all bodies in simulation world frame. Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to @@ -832,7 +834,7 @@ def body_link_pos_w(self) -> wp.array: @property @abstractmethod - def body_link_quat_w(self) -> wp.array: + def body_link_quat_w(self) -> ProxyArray: """Orientation (x, y, z, w) of all bodies in simulation world frame. Shape is (num_instances, num_bodies), dtype = wp.quatf. In torch this resolves to @@ -844,7 +846,7 @@ def body_link_quat_w(self) -> wp.array: @property @abstractmethod - def body_link_lin_vel_w(self) -> wp.array: + def body_link_lin_vel_w(self) -> ProxyArray: """Linear velocity of all bodies in simulation world frame. Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to @@ -856,7 +858,7 @@ def body_link_lin_vel_w(self) -> wp.array: @property @abstractmethod - def body_link_ang_vel_w(self) -> wp.array: + def body_link_ang_vel_w(self) -> ProxyArray: """Angular velocity of all bodies in simulation world frame. Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to @@ -868,7 +870,7 @@ def body_link_ang_vel_w(self) -> wp.array: @property @abstractmethod - def body_com_pos_w(self) -> wp.array: + def body_com_pos_w(self) -> ProxyArray: """Positions of all bodies in simulation world frame. Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to @@ -880,7 +882,7 @@ def body_com_pos_w(self) -> wp.array: @property @abstractmethod - def body_com_quat_w(self) -> wp.array: + def body_com_quat_w(self) -> ProxyArray: """Orientation (x, y, z, w) of the principal axes of inertia of all bodies in simulation world frame. Shape is (num_instances, num_bodies), dtype = wp.quatf. In torch this resolves to @@ -892,7 +894,7 @@ def body_com_quat_w(self) -> wp.array: @property @abstractmethod - def body_com_lin_vel_w(self) -> wp.array: + def body_com_lin_vel_w(self) -> ProxyArray: """Linear velocity of all bodies in simulation world frame. Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to @@ -904,7 +906,7 @@ def body_com_lin_vel_w(self) -> wp.array: @property @abstractmethod - def body_com_ang_vel_w(self) -> wp.array: + def body_com_ang_vel_w(self) -> ProxyArray: """Angular velocity of all bodies in simulation world frame. Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to @@ -916,7 +918,7 @@ def body_com_ang_vel_w(self) -> wp.array: @property @abstractmethod - def body_com_lin_acc_w(self) -> wp.array: + def body_com_lin_acc_w(self) -> ProxyArray: """Linear acceleration of all bodies in simulation world frame. Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to @@ -928,7 +930,7 @@ def body_com_lin_acc_w(self) -> wp.array: @property @abstractmethod - def body_com_ang_acc_w(self) -> wp.array: + def body_com_ang_acc_w(self) -> ProxyArray: """Angular acceleration of all bodies in simulation world frame. Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to @@ -940,7 +942,7 @@ def body_com_ang_acc_w(self) -> wp.array: @property @abstractmethod - def body_com_pos_b(self) -> wp.array: + def body_com_pos_b(self) -> ProxyArray: """Center of mass position of all of the bodies in their respective link frames. Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to @@ -952,7 +954,7 @@ def body_com_pos_b(self) -> wp.array: @property @abstractmethod - def body_com_quat_b(self) -> wp.array: + def body_com_quat_b(self) -> ProxyArray: """Orientation (x, y, z, w) of the principal axes of inertia of all of the bodies in their respective link frames. @@ -989,122 +991,122 @@ def _create_buffers(self) -> None: """ @property - def root_pose_w(self) -> wp.array: + def root_pose_w(self) -> ProxyArray: """Shorthand for :attr:`root_link_pose_w`.""" return self.root_link_pose_w @property - def root_pos_w(self) -> wp.array: + def root_pos_w(self) -> ProxyArray: """Shorthand for :attr:`root_link_pos_w`.""" return self.root_link_pos_w @property - def root_quat_w(self) -> wp.array: + def root_quat_w(self) -> ProxyArray: """Shorthand for :attr:`root_link_quat_w`.""" return self.root_link_quat_w @property - def root_vel_w(self) -> wp.array: + def root_vel_w(self) -> ProxyArray: """Shorthand for :attr:`root_com_vel_w`.""" return self.root_com_vel_w @property - def root_lin_vel_w(self) -> wp.array: + def root_lin_vel_w(self) -> ProxyArray: """Shorthand for :attr:`root_com_lin_vel_w`.""" return self.root_com_lin_vel_w @property - def root_ang_vel_w(self) -> wp.array: + def root_ang_vel_w(self) -> ProxyArray: """Shorthand for :attr:`root_com_ang_vel_w`.""" return self.root_com_ang_vel_w @property - def root_lin_vel_b(self) -> wp.array: + def root_lin_vel_b(self) -> ProxyArray: """Shorthand for :attr:`root_com_lin_vel_b`.""" return self.root_com_lin_vel_b @property - def root_ang_vel_b(self) -> wp.array: + def root_ang_vel_b(self) -> ProxyArray: """Shorthand for :attr:`root_com_ang_vel_b`.""" return self.root_com_ang_vel_b @property - def body_pose_w(self) -> wp.array: + def body_pose_w(self) -> ProxyArray: """Shorthand for :attr:`body_link_pose_w`.""" return self.body_link_pose_w @property - def body_pos_w(self) -> wp.array: + def body_pos_w(self) -> ProxyArray: """Shorthand for :attr:`body_link_pos_w`.""" return self.body_link_pos_w @property - def body_quat_w(self) -> wp.array: + def body_quat_w(self) -> ProxyArray: """Shorthand for :attr:`body_link_quat_w`.""" return self.body_link_quat_w @property - def body_vel_w(self) -> wp.array: + def body_vel_w(self) -> ProxyArray: """Shorthand for :attr:`body_com_vel_w`.""" return self.body_com_vel_w @property - def body_lin_vel_w(self) -> wp.array: + def body_lin_vel_w(self) -> ProxyArray: """Shorthand for :attr:`body_com_lin_vel_w`.""" return self.body_com_lin_vel_w @property - def body_ang_vel_w(self) -> wp.array: + def body_ang_vel_w(self) -> ProxyArray: """Shorthand for :attr:`body_com_ang_vel_w`.""" return self.body_com_ang_vel_w @property - def body_acc_w(self) -> wp.array: + def body_acc_w(self) -> ProxyArray: """Shorthand for :attr:`body_com_acc_w`.""" return self.body_com_acc_w @property - def body_lin_acc_w(self) -> wp.array: + def body_lin_acc_w(self) -> ProxyArray: """Shorthand for :attr:`body_com_lin_acc_w`.""" return self.body_com_lin_acc_w @property - def body_ang_acc_w(self) -> wp.array: + def body_ang_acc_w(self) -> ProxyArray: """Shorthand for :attr:`body_com_ang_acc_w`.""" return self.body_com_ang_acc_w @property - def com_pos_b(self) -> wp.array: + def com_pos_b(self) -> ProxyArray: """Shorthand for :attr:`body_com_pos_b`.""" return self.body_com_pos_b @property - def com_quat_b(self) -> wp.array: + def com_quat_b(self) -> ProxyArray: """Shorthand for :attr:`body_com_quat_b`.""" return self.body_com_quat_b @property - def joint_limits(self) -> wp.array: + def joint_limits(self) -> ProxyArray: """Shorthand for :attr:`joint_pos_limits`.""" return self.joint_pos_limits @property - def default_joint_limits(self) -> wp.array: + def default_joint_limits(self) -> ProxyArray: """Shorthand for :attr:`default_joint_pos_limits`.""" return self.default_joint_pos_limits @property - def joint_velocity_limits(self) -> wp.array: + def joint_velocity_limits(self) -> ProxyArray: """Shorthand for :attr:`joint_vel_limits`.""" return self.joint_vel_limits @property - def joint_friction(self) -> wp.array: + def joint_friction(self) -> ProxyArray: """Shorthand for :attr:`joint_friction_coeff`.""" return self.joint_friction_coeff @property - def fixed_tendon_limit(self) -> wp.array: + def fixed_tendon_limit(self) -> ProxyArray: """Shorthand for :attr:`fixed_tendon_pos_limits`.""" return self.fixed_tendon_pos_limits @@ -1113,7 +1115,7 @@ def fixed_tendon_limit(self) -> wp.array: """ @property - def default_mass(self) -> wp.array: + def default_mass(self) -> ProxyArray: """Deprecated property. Please use :attr:`body_mass` instead and manage the default mass manually.""" warnings.warn( "The `default_mass` property will be deprecated in a IsaacLab 4.0. Please use `body_mass` instead. " @@ -1122,11 +1124,11 @@ def default_mass(self) -> wp.array: stacklevel=2, ) if self._default_mass is None: - self._default_mass = wp.clone(self.body_mass, self.device) - return self._default_mass + self._default_mass = wp.clone(self.body_mass.warp, self.device) + return ProxyArray(self._default_mass) @property - def default_inertia(self) -> wp.array: + def default_inertia(self) -> ProxyArray: """Deprecated property. Please use :attr:`body_inertia` instead and manage the default inertia manually.""" warnings.warn( "The `default_inertia` property will be deprecated in a IsaacLab 4.0. Please use `body_inertia` instead. " @@ -1135,11 +1137,11 @@ def default_inertia(self) -> wp.array: stacklevel=2, ) if self._default_inertia is None: - self._default_inertia = wp.clone(self.body_inertia, self.device) - return self._default_inertia + self._default_inertia = wp.clone(self.body_inertia.warp, self.device) + return ProxyArray(self._default_inertia) @property - def default_joint_stiffness(self) -> wp.array: + def default_joint_stiffness(self) -> ProxyArray: """Deprecated property. Please use :attr:`joint_stiffness` instead and manage the default joint stiffness manually.""" warnings.warn( @@ -1149,11 +1151,11 @@ def default_joint_stiffness(self) -> wp.array: stacklevel=2, ) if self._default_joint_stiffness is None: - self._default_joint_stiffness = wp.clone(self.joint_stiffness, self.device) - return self._default_joint_stiffness + self._default_joint_stiffness = wp.clone(self.joint_stiffness.warp, self.device) + return ProxyArray(self._default_joint_stiffness) @property - def default_joint_damping(self) -> wp.array: + def default_joint_damping(self) -> ProxyArray: """Deprecated property. Please use :attr:`joint_damping` instead and manage the default joint damping manually.""" warnings.warn( @@ -1163,11 +1165,11 @@ def default_joint_damping(self) -> wp.array: stacklevel=2, ) if self._default_joint_damping is None: - self._default_joint_damping = wp.clone(self.joint_damping, self.device) - return self._default_joint_damping + self._default_joint_damping = wp.clone(self.joint_damping.warp, self.device) + return ProxyArray(self._default_joint_damping) @property - def default_joint_armature(self) -> wp.array: + def default_joint_armature(self) -> ProxyArray: """Deprecated property. Please use :attr:`joint_armature` instead and manage the default joint armature manually.""" warnings.warn( @@ -1177,11 +1179,11 @@ def default_joint_armature(self) -> wp.array: stacklevel=2, ) if self._default_joint_armature is None: - self._default_joint_armature = wp.clone(self.joint_armature, self.device) - return self._default_joint_armature + self._default_joint_armature = wp.clone(self.joint_armature.warp, self.device) + return ProxyArray(self._default_joint_armature) @property - def default_joint_friction_coeff(self) -> wp.array: + def default_joint_friction_coeff(self) -> ProxyArray: """Deprecated property. Please use :attr:`joint_friction_coeff` instead and manage the default joint friction coefficient manually.""" warnings.warn( @@ -1191,11 +1193,11 @@ def default_joint_friction_coeff(self) -> wp.array: stacklevel=2, ) if self._default_joint_friction_coeff is None: - self._default_joint_friction_coeff = wp.clone(self.joint_friction_coeff, self.device) - return self._default_joint_friction_coeff + self._default_joint_friction_coeff = wp.clone(self.joint_friction_coeff.warp, self.device) + return ProxyArray(self._default_joint_friction_coeff) @property - def default_joint_viscous_friction_coeff(self) -> wp.array: + def default_joint_viscous_friction_coeff(self) -> ProxyArray: """Deprecated property. Please use :attr:`joint_viscous_friction_coeff` instead and manage the default joint viscous friction coefficient manually.""" warnings.warn( @@ -1205,11 +1207,11 @@ def default_joint_viscous_friction_coeff(self) -> wp.array: stacklevel=2, ) if self._default_joint_viscous_friction_coeff is None: - self._default_joint_viscous_friction_coeff = wp.clone(self.joint_viscous_friction_coeff, self.device) - return self._default_joint_viscous_friction_coeff + self._default_joint_viscous_friction_coeff = wp.clone(self.joint_viscous_friction_coeff.warp, self.device) + return ProxyArray(self._default_joint_viscous_friction_coeff) @property - def default_joint_pos_limits(self) -> wp.array: + def default_joint_pos_limits(self) -> ProxyArray: """Deprecated property. Please use :attr:`joint_pos_limits` instead and manage the default joint position limits manually.""" warnings.warn( @@ -1219,11 +1221,11 @@ def default_joint_pos_limits(self) -> wp.array: stacklevel=2, ) if self._default_joint_pos_limits is None: - self._default_joint_pos_limits = wp.clone(self.joint_pos_limits, self.device) - return self._default_joint_pos_limits + self._default_joint_pos_limits = wp.clone(self.joint_pos_limits.warp, self.device) + return ProxyArray(self._default_joint_pos_limits) @property - def default_fixed_tendon_stiffness(self) -> wp.array: + def default_fixed_tendon_stiffness(self) -> ProxyArray: """Deprecated property. Please use :attr:`fixed_tendon_stiffness` instead and manage the default fixed tendon stiffness manually.""" warnings.warn( @@ -1233,11 +1235,11 @@ def default_fixed_tendon_stiffness(self) -> wp.array: stacklevel=2, ) if self._default_fixed_tendon_stiffness is None: - self._default_fixed_tendon_stiffness = wp.clone(self.fixed_tendon_stiffness, self.device) - return self._default_fixed_tendon_stiffness + self._default_fixed_tendon_stiffness = wp.clone(self.fixed_tendon_stiffness.warp, self.device) + return ProxyArray(self._default_fixed_tendon_stiffness) @property - def default_fixed_tendon_damping(self) -> wp.array: + def default_fixed_tendon_damping(self) -> ProxyArray: """Deprecated property. Please use :attr:`fixed_tendon_damping` instead and manage the default fixed tendon damping manually.""" warnings.warn( @@ -1247,11 +1249,11 @@ def default_fixed_tendon_damping(self) -> wp.array: stacklevel=2, ) if self._default_fixed_tendon_damping is None: - self._default_fixed_tendon_damping = wp.clone(self.fixed_tendon_damping, self.device) - return self._default_fixed_tendon_damping + self._default_fixed_tendon_damping = wp.clone(self.fixed_tendon_damping.warp, self.device) + return ProxyArray(self._default_fixed_tendon_damping) @property - def default_fixed_tendon_limit_stiffness(self) -> wp.array: + def default_fixed_tendon_limit_stiffness(self) -> ProxyArray: """Deprecated property. Please use :attr:`fixed_tendon_limit_stiffness` instead and manage the default fixed tendon limit stiffness manually.""" warnings.warn( @@ -1261,11 +1263,11 @@ def default_fixed_tendon_limit_stiffness(self) -> wp.array: stacklevel=2, ) if self._default_fixed_tendon_limit_stiffness is None: - self._default_fixed_tendon_limit_stiffness = wp.clone(self.fixed_tendon_limit_stiffness, self.device) - return self._default_fixed_tendon_limit_stiffness + self._default_fixed_tendon_limit_stiffness = wp.clone(self.fixed_tendon_limit_stiffness.warp, self.device) + return ProxyArray(self._default_fixed_tendon_limit_stiffness) @property - def default_fixed_tendon_rest_length(self) -> wp.array: + def default_fixed_tendon_rest_length(self) -> ProxyArray: """Deprecated property. Please use :attr:`fixed_tendon_rest_length` instead and manage the default fixed tendon rest length manually.""" warnings.warn( @@ -1275,11 +1277,11 @@ def default_fixed_tendon_rest_length(self) -> wp.array: stacklevel=2, ) if self._default_fixed_tendon_rest_length is None: - self._default_fixed_tendon_rest_length = wp.clone(self.fixed_tendon_rest_length, self.device) - return self._default_fixed_tendon_rest_length + self._default_fixed_tendon_rest_length = wp.clone(self.fixed_tendon_rest_length.warp, self.device) + return ProxyArray(self._default_fixed_tendon_rest_length) @property - def default_fixed_tendon_offset(self) -> wp.array: + def default_fixed_tendon_offset(self) -> ProxyArray: """Deprecated property. Please use :attr:`fixed_tendon_offset` instead and manage the default fixed tendon offset manually.""" warnings.warn( @@ -1289,11 +1291,11 @@ def default_fixed_tendon_offset(self) -> wp.array: stacklevel=2, ) if self._default_fixed_tendon_offset is None: - self._default_fixed_tendon_offset = wp.clone(self.fixed_tendon_offset, self.device) - return self._default_fixed_tendon_offset + self._default_fixed_tendon_offset = wp.clone(self.fixed_tendon_offset.warp, self.device) + return ProxyArray(self._default_fixed_tendon_offset) @property - def default_fixed_tendon_pos_limits(self) -> wp.array: + def default_fixed_tendon_pos_limits(self) -> ProxyArray: """Deprecated property. Please use :attr:`fixed_tendon_pos_limits` instead and manage the default fixed tendon position limits manually.""" warnings.warn( @@ -1303,11 +1305,11 @@ def default_fixed_tendon_pos_limits(self) -> wp.array: stacklevel=2, ) if self._default_fixed_tendon_pos_limits is None: - self._default_fixed_tendon_pos_limits = wp.clone(self.fixed_tendon_pos_limits, self.device) - return self._default_fixed_tendon_pos_limits + self._default_fixed_tendon_pos_limits = wp.clone(self.fixed_tendon_pos_limits.warp, self.device) + return ProxyArray(self._default_fixed_tendon_pos_limits) @property - def default_spatial_tendon_stiffness(self) -> wp.array: + def default_spatial_tendon_stiffness(self) -> ProxyArray: """Deprecated property. Please use :attr:`spatial_tendon_stiffness` instead and manage the default spatial tendon stiffness manually.""" warnings.warn( @@ -1317,11 +1319,11 @@ def default_spatial_tendon_stiffness(self) -> wp.array: stacklevel=2, ) if self._default_spatial_tendon_stiffness is None: - self._default_spatial_tendon_stiffness = wp.clone(self.spatial_tendon_stiffness, self.device) - return self._default_spatial_tendon_stiffness + self._default_spatial_tendon_stiffness = wp.clone(self.spatial_tendon_stiffness.warp, self.device) + return ProxyArray(self._default_spatial_tendon_stiffness) @property - def default_spatial_tendon_damping(self) -> wp.array: + def default_spatial_tendon_damping(self) -> ProxyArray: """Deprecated property. Please use :attr:`spatial_tendon_damping` instead and manage the default spatial tendon damping manually.""" warnings.warn( @@ -1331,11 +1333,11 @@ def default_spatial_tendon_damping(self) -> wp.array: stacklevel=2, ) if self._default_spatial_tendon_damping is None: - self._default_spatial_tendon_damping = wp.clone(self.spatial_tendon_damping, self.device) - return self._default_spatial_tendon_damping + self._default_spatial_tendon_damping = wp.clone(self.spatial_tendon_damping.warp, self.device) + return ProxyArray(self._default_spatial_tendon_damping) @property - def default_spatial_tendon_limit_stiffness(self) -> wp.array: + def default_spatial_tendon_limit_stiffness(self) -> ProxyArray: """Deprecated property. Please use :attr:`spatial_tendon_limit_stiffness` instead and manage the default spatial tendon limit stiffness manually.""" warnings.warn( @@ -1345,11 +1347,13 @@ def default_spatial_tendon_limit_stiffness(self) -> wp.array: stacklevel=2, ) if self._default_spatial_tendon_limit_stiffness is None: - self._default_spatial_tendon_limit_stiffness = wp.clone(self.spatial_tendon_limit_stiffness, self.device) - return self._default_spatial_tendon_limit_stiffness + self._default_spatial_tendon_limit_stiffness = wp.clone( + self.spatial_tendon_limit_stiffness.warp, self.device + ) + return ProxyArray(self._default_spatial_tendon_limit_stiffness) @property - def default_spatial_tendon_offset(self) -> wp.array: + def default_spatial_tendon_offset(self) -> ProxyArray: """Deprecated property. Please use :attr:`spatial_tendon_offset` instead and manage the default spatial tendon offset manually.""" warnings.warn( @@ -1359,11 +1363,11 @@ def default_spatial_tendon_offset(self) -> wp.array: stacklevel=2, ) if self._default_spatial_tendon_offset is None: - self._default_spatial_tendon_offset = wp.clone(self.spatial_tendon_offset, self.device) - return self._default_spatial_tendon_offset + self._default_spatial_tendon_offset = wp.clone(self.spatial_tendon_offset.warp, self.device) + return ProxyArray(self._default_spatial_tendon_offset) @property - def default_fixed_tendon_limit(self) -> wp.array: + def default_fixed_tendon_limit(self) -> ProxyArray: """Deprecated property. Please use :attr:`default_fixed_tendon_pos_limits` instead.""" warnings.warn( "The `default_fixed_tendon_limit` property will be deprecated in a IsaacLab 4.0. Please use" @@ -1374,7 +1378,7 @@ def default_fixed_tendon_limit(self) -> wp.array: return self.default_fixed_tendon_pos_limits @property - def default_joint_friction(self) -> wp.array: + def default_joint_friction(self) -> ProxyArray: """Deprecated property. Please use :attr:`default_joint_friction_coeff` instead.""" warnings.warn( "The `default_joint_friction` property will be deprecated in a IsaacLab 4.0. Please use" diff --git a/source/isaaclab/isaaclab/assets/rigid_object/base_rigid_object_data.py b/source/isaaclab/isaaclab/assets/rigid_object/base_rigid_object_data.py index e04a1ffa8bc1..e73302db6436 100644 --- a/source/isaaclab/isaaclab/assets/rigid_object/base_rigid_object_data.py +++ b/source/isaaclab/isaaclab/assets/rigid_object/base_rigid_object_data.py @@ -8,6 +8,8 @@ import warp as wp +from isaaclab.utils.warp import ProxyArray + class BaseRigidObjectData(ABC): """Data container for a rigid object. @@ -61,7 +63,7 @@ def update(self, dt: float) -> None: @property @abstractmethod - def default_root_pose(self) -> wp.array: + def default_root_pose(self) -> ProxyArray: """Default root pose ``[pos, quat]`` in local environment frame. The position and quaternion are of the rigid body's actor frame. @@ -71,7 +73,7 @@ def default_root_pose(self) -> wp.array: @property @abstractmethod - def default_root_vel(self) -> wp.array: + def default_root_vel(self) -> ProxyArray: """Default root velocity ``[lin_vel, ang_vel]`` in local environment frame. The linear and angular velocities are of the rigid body's center of mass frame. @@ -81,7 +83,7 @@ def default_root_vel(self) -> wp.array: @property @abstractmethod - def default_root_state(self) -> wp.array: + def default_root_state(self) -> ProxyArray: """Deprecated, same as :attr:`default_root_pose` and :attr:`default_root_vel`.""" raise NotImplementedError() @@ -91,7 +93,7 @@ def default_root_state(self) -> wp.array: @property @abstractmethod - def root_link_pose_w(self) -> wp.array: + def root_link_pose_w(self) -> ProxyArray: """Root link pose ``[pos, quat]`` in simulation world frame. Shape is (num_instances,), dtype = wp.transformf. In torch this resolves to (num_instances, 7). @@ -103,7 +105,7 @@ def root_link_pose_w(self) -> wp.array: @property @abstractmethod - def root_link_vel_w(self) -> wp.array: + def root_link_vel_w(self) -> ProxyArray: """Root link velocity ``[lin_vel, ang_vel]`` in simulation world frame. Shape is (num_instances,), dtype = wp.spatial_vectorf. In torch this resolves to (num_instances, 6). @@ -115,7 +117,7 @@ def root_link_vel_w(self) -> wp.array: @property @abstractmethod - def root_com_pose_w(self) -> wp.array: + def root_com_pose_w(self) -> ProxyArray: """Root center of mass pose ``[pos, quat]`` in simulation world frame. Shape is (num_instances,), dtype = wp.transformf. In torch this resolves to (num_instances, 7). @@ -127,7 +129,7 @@ def root_com_pose_w(self) -> wp.array: @property @abstractmethod - def root_com_vel_w(self) -> wp.array: + def root_com_vel_w(self) -> ProxyArray: """Root center of mass velocity ``[lin_vel, ang_vel]`` in simulation world frame. Shape is (num_instances,), dtype = wp.spatial_vectorf. In torch this resolves to (num_instances, 6). @@ -139,19 +141,19 @@ def root_com_vel_w(self) -> wp.array: @property @abstractmethod - def root_state_w(self) -> wp.array: + def root_state_w(self) -> ProxyArray: """Deprecated, same as :attr:`root_link_pose_w` and :attr:`root_com_vel_w`.""" raise NotImplementedError() @property @abstractmethod - def root_link_state_w(self) -> wp.array: + def root_link_state_w(self) -> ProxyArray: """Deprecated, same as :attr:`root_link_pose_w` and :attr:`root_link_vel_w`.""" raise NotImplementedError() @property @abstractmethod - def root_com_state_w(self) -> wp.array: + def root_com_state_w(self) -> ProxyArray: """Deprecated, same as :attr:`root_com_pose_w` and :attr:`root_com_vel_w`.""" raise NotImplementedError() @@ -161,7 +163,7 @@ def root_com_state_w(self) -> wp.array: @property @abstractmethod - def body_link_pose_w(self) -> wp.array: + def body_link_pose_w(self) -> ProxyArray: """Body link pose ``[pos, quat]`` in simulation world frame. Shape is (num_instances, 1), dtype = wp.transformf. @@ -174,7 +176,7 @@ def body_link_pose_w(self) -> wp.array: @property @abstractmethod - def body_link_vel_w(self) -> wp.array: + def body_link_vel_w(self) -> ProxyArray: """Body link velocity ``[lin_vel, ang_vel]`` in simulation world frame. Shape is (num_instances, 1), dtype = wp.spatial_vectorf. @@ -187,7 +189,7 @@ def body_link_vel_w(self) -> wp.array: @property @abstractmethod - def body_com_pose_w(self) -> wp.array: + def body_com_pose_w(self) -> ProxyArray: """Body center of mass pose ``[pos, quat]`` in simulation world frame. Shape is (num_instances, 1), dtype = wp.transformf. @@ -200,7 +202,7 @@ def body_com_pose_w(self) -> wp.array: @property @abstractmethod - def body_com_vel_w(self) -> wp.array: + def body_com_vel_w(self) -> ProxyArray: """Body center of mass velocity ``[lin_vel, ang_vel]`` in simulation world frame. Shape is (num_instances, 1), dtype = wp.spatial_vectorf. @@ -213,25 +215,25 @@ def body_com_vel_w(self) -> wp.array: @property @abstractmethod - def body_state_w(self) -> wp.array: + def body_state_w(self) -> ProxyArray: """Deprecated, same as :attr:`body_link_pose_w` and :attr:`body_com_vel_w`.""" raise NotImplementedError() @property @abstractmethod - def body_link_state_w(self) -> wp.array: + def body_link_state_w(self) -> ProxyArray: """Deprecated, same as :attr:`body_link_pose_w` and :attr:`body_link_vel_w`.""" raise NotImplementedError() @property @abstractmethod - def body_com_state_w(self) -> wp.array: + def body_com_state_w(self) -> ProxyArray: """Deprecated, same as :attr:`body_com_pose_w` and :attr:`body_com_vel_w`.""" raise NotImplementedError() @property @abstractmethod - def body_com_acc_w(self) -> wp.array: + def body_com_acc_w(self) -> ProxyArray: """Acceleration of all bodies ``[lin_acc, ang_acc]`` in the simulation world frame. Shape is (num_instances, 1), dtype = wp.spatial_vectorf. @@ -243,7 +245,7 @@ def body_com_acc_w(self) -> wp.array: @property @abstractmethod - def body_com_pose_b(self) -> wp.array: + def body_com_pose_b(self) -> ProxyArray: """Center of mass pose ``[pos, quat]`` of all bodies in their respective body's link frames. Shape is (num_instances, 1), dtype = wp.transformf. @@ -256,7 +258,7 @@ def body_com_pose_b(self) -> wp.array: @property @abstractmethod - def body_mass(self) -> wp.array: + def body_mass(self) -> ProxyArray: """Mass of all bodies in the simulation world frame. Shape is (num_instances, 1), dtype = wp.float32. @@ -266,7 +268,7 @@ def body_mass(self) -> wp.array: @property @abstractmethod - def body_inertia(self) -> wp.array: + def body_inertia(self) -> ProxyArray: """Inertia of all bodies in the simulation world frame. Shape is (num_instances, 1, 9), dtype = wp.float32. @@ -280,7 +282,7 @@ def body_inertia(self) -> wp.array: @property @abstractmethod - def projected_gravity_b(self) -> wp.array: + def projected_gravity_b(self) -> ProxyArray: """Projection of the gravity direction on base frame. Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). @@ -289,7 +291,7 @@ def projected_gravity_b(self) -> wp.array: @property @abstractmethod - def heading_w(self) -> wp.array: + def heading_w(self) -> ProxyArray: """Yaw heading of the base frame (in radians). Shape is (num_instances,), dtype = wp.float32. In torch this resolves to (num_instances,). @@ -302,7 +304,7 @@ def heading_w(self) -> wp.array: @property @abstractmethod - def root_link_lin_vel_b(self) -> wp.array: + def root_link_lin_vel_b(self) -> ProxyArray: """Root link linear velocity in base frame. Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). @@ -314,7 +316,7 @@ def root_link_lin_vel_b(self) -> wp.array: @property @abstractmethod - def root_link_ang_vel_b(self) -> wp.array: + def root_link_ang_vel_b(self) -> ProxyArray: """Root link angular velocity in base frame. Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). @@ -326,7 +328,7 @@ def root_link_ang_vel_b(self) -> wp.array: @property @abstractmethod - def root_com_lin_vel_b(self) -> wp.array: + def root_com_lin_vel_b(self) -> ProxyArray: """Root center of mass linear velocity in base frame. Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). @@ -338,7 +340,7 @@ def root_com_lin_vel_b(self) -> wp.array: @property @abstractmethod - def root_com_ang_vel_b(self) -> wp.array: + def root_com_ang_vel_b(self) -> ProxyArray: """Root center of mass angular velocity in base frame. Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). @@ -354,7 +356,7 @@ def root_com_ang_vel_b(self) -> wp.array: @property @abstractmethod - def root_link_pos_w(self) -> wp.array: + def root_link_pos_w(self) -> ProxyArray: """Root link position in simulation world frame. Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). @@ -365,7 +367,7 @@ def root_link_pos_w(self) -> wp.array: @property @abstractmethod - def root_link_quat_w(self) -> wp.array: + def root_link_quat_w(self) -> ProxyArray: """Root link orientation (x, y, z, w) in simulation world frame. Shape is (num_instances,), dtype = wp.quatf. In torch this resolves to (num_instances, 4). @@ -376,7 +378,7 @@ def root_link_quat_w(self) -> wp.array: @property @abstractmethod - def root_link_lin_vel_w(self) -> wp.array: + def root_link_lin_vel_w(self) -> ProxyArray: """Root linear velocity in simulation world frame. Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). @@ -387,7 +389,7 @@ def root_link_lin_vel_w(self) -> wp.array: @property @abstractmethod - def root_link_ang_vel_w(self) -> wp.array: + def root_link_ang_vel_w(self) -> ProxyArray: """Root link angular velocity in simulation world frame. Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). @@ -398,7 +400,7 @@ def root_link_ang_vel_w(self) -> wp.array: @property @abstractmethod - def root_com_pos_w(self) -> wp.array: + def root_com_pos_w(self) -> ProxyArray: """Root center of mass position in simulation world frame. Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). @@ -409,7 +411,7 @@ def root_com_pos_w(self) -> wp.array: @property @abstractmethod - def root_com_quat_w(self) -> wp.array: + def root_com_quat_w(self) -> ProxyArray: """Root center of mass orientation (x, y, z, w) in simulation world frame. Shape is (num_instances,), dtype = wp.quatf. In torch this resolves to (num_instances, 4). @@ -420,7 +422,7 @@ def root_com_quat_w(self) -> wp.array: @property @abstractmethod - def root_com_lin_vel_w(self) -> wp.array: + def root_com_lin_vel_w(self) -> ProxyArray: """Root center of mass linear velocity in simulation world frame. Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). @@ -431,7 +433,7 @@ def root_com_lin_vel_w(self) -> wp.array: @property @abstractmethod - def root_com_ang_vel_w(self) -> wp.array: + def root_com_ang_vel_w(self) -> ProxyArray: """Root center of mass angular velocity in simulation world frame. Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). @@ -442,7 +444,7 @@ def root_com_ang_vel_w(self) -> wp.array: @property @abstractmethod - def body_link_pos_w(self) -> wp.array: + def body_link_pos_w(self) -> ProxyArray: """Positions of all bodies in simulation world frame. Shape is (num_instances, 1), dtype = wp.vec3f. In torch this resolves to (num_instances, 1, 3). @@ -453,7 +455,7 @@ def body_link_pos_w(self) -> wp.array: @property @abstractmethod - def body_link_quat_w(self) -> wp.array: + def body_link_quat_w(self) -> ProxyArray: """Orientation (x, y, z, w) of all bodies in simulation world frame. Shape is (num_instances, 1), dtype = wp.quatf. In torch this resolves to (num_instances, 1, 4). @@ -464,7 +466,7 @@ def body_link_quat_w(self) -> wp.array: @property @abstractmethod - def body_link_lin_vel_w(self) -> wp.array: + def body_link_lin_vel_w(self) -> ProxyArray: """Linear velocity of all bodies in simulation world frame. Shape is (num_instances, 1), dtype = wp.vec3f. In torch this resolves to (num_instances, 1, 3). @@ -475,7 +477,7 @@ def body_link_lin_vel_w(self) -> wp.array: @property @abstractmethod - def body_link_ang_vel_w(self) -> wp.array: + def body_link_ang_vel_w(self) -> ProxyArray: """Angular velocity of all bodies in simulation world frame. Shape is (num_instances, 1), dtype = wp.vec3f. In torch this resolves to (num_instances, 1, 3). @@ -486,7 +488,7 @@ def body_link_ang_vel_w(self) -> wp.array: @property @abstractmethod - def body_com_pos_w(self) -> wp.array: + def body_com_pos_w(self) -> ProxyArray: """Positions of all bodies' center of mass in simulation world frame. Shape is (num_instances, 1), dtype = wp.vec3f. In torch this resolves to (num_instances, 1, 3). @@ -497,7 +499,7 @@ def body_com_pos_w(self) -> wp.array: @property @abstractmethod - def body_com_quat_w(self) -> wp.array: + def body_com_quat_w(self) -> ProxyArray: """Orientation (x, y, z, w) of the principal axes of inertia of all bodies in simulation world frame. Shape is (num_instances, 1), dtype = wp.quatf. In torch this resolves to (num_instances, 1, 4). @@ -508,7 +510,7 @@ def body_com_quat_w(self) -> wp.array: @property @abstractmethod - def body_com_lin_vel_w(self) -> wp.array: + def body_com_lin_vel_w(self) -> ProxyArray: """Linear velocity of all bodies in simulation world frame. Shape is (num_instances, 1), dtype = wp.vec3f. In torch this resolves to (num_instances, 1, 3). @@ -519,7 +521,7 @@ def body_com_lin_vel_w(self) -> wp.array: @property @abstractmethod - def body_com_ang_vel_w(self) -> wp.array: + def body_com_ang_vel_w(self) -> ProxyArray: """Angular velocity of all bodies in simulation world frame. Shape is (num_instances, 1), dtype = wp.vec3f. In torch this resolves to (num_instances, 1, 3). @@ -530,7 +532,7 @@ def body_com_ang_vel_w(self) -> wp.array: @property @abstractmethod - def body_com_lin_acc_w(self) -> wp.array: + def body_com_lin_acc_w(self) -> ProxyArray: """Linear acceleration of all bodies in simulation world frame. Shape is (num_instances, 1), dtype = wp.vec3f. In torch this resolves to (num_instances, 1, 3). @@ -541,7 +543,7 @@ def body_com_lin_acc_w(self) -> wp.array: @property @abstractmethod - def body_com_ang_acc_w(self) -> wp.array: + def body_com_ang_acc_w(self) -> ProxyArray: """Angular acceleration of all bodies in simulation world frame. Shape is (num_instances, 1), dtype = wp.vec3f. In torch this resolves to (num_instances, 1, 3). @@ -552,7 +554,7 @@ def body_com_ang_acc_w(self) -> wp.array: @property @abstractmethod - def body_com_pos_b(self) -> wp.array: + def body_com_pos_b(self) -> ProxyArray: """Center of mass position of all of the bodies in their respective link frames. Shape is (num_instances, 1), dtype = wp.vec3f. In torch this resolves to (num_instances, 1, 3). @@ -563,7 +565,7 @@ def body_com_pos_b(self) -> wp.array: @property @abstractmethod - def body_com_quat_b(self) -> wp.array: + def body_com_quat_b(self) -> ProxyArray: """Orientation (x, y, z, w) of the principal axes of inertia of all of the bodies in their respective link frames. @@ -583,97 +585,97 @@ def _create_buffers(self) -> None: """ @property - def root_pose_w(self) -> wp.array: + def root_pose_w(self) -> ProxyArray: """Shorthand for :attr:`root_link_pose_w`.""" return self.root_link_pose_w @property - def root_pos_w(self) -> wp.array: + def root_pos_w(self) -> ProxyArray: """Shorthand for :attr:`root_link_pos_w`.""" return self.root_link_pos_w @property - def root_quat_w(self) -> wp.array: + def root_quat_w(self) -> ProxyArray: """Shorthand for :attr:`root_link_quat_w`.""" return self.root_link_quat_w @property - def root_vel_w(self) -> wp.array: + def root_vel_w(self) -> ProxyArray: """Shorthand for :attr:`root_com_vel_w`.""" return self.root_com_vel_w @property - def root_lin_vel_w(self) -> wp.array: + def root_lin_vel_w(self) -> ProxyArray: """Shorthand for :attr:`root_com_lin_vel_w`.""" return self.root_com_lin_vel_w @property - def root_ang_vel_w(self) -> wp.array: + def root_ang_vel_w(self) -> ProxyArray: """Shorthand for :attr:`root_com_ang_vel_w`.""" return self.root_com_ang_vel_w @property - def root_lin_vel_b(self) -> wp.array: + def root_lin_vel_b(self) -> ProxyArray: """Shorthand for :attr:`root_com_lin_vel_b`.""" return self.root_com_lin_vel_b @property - def root_ang_vel_b(self) -> wp.array: + def root_ang_vel_b(self) -> ProxyArray: """Shorthand for :attr:`root_com_ang_vel_b`.""" return self.root_com_ang_vel_b @property - def body_pose_w(self) -> wp.array: + def body_pose_w(self) -> ProxyArray: """Shorthand for :attr:`body_link_pose_w`.""" return self.body_link_pose_w @property - def body_pos_w(self) -> wp.array: + def body_pos_w(self) -> ProxyArray: """Shorthand for :attr:`body_link_pos_w`.""" return self.body_link_pos_w @property - def body_quat_w(self) -> wp.array: + def body_quat_w(self) -> ProxyArray: """Shorthand for :attr:`body_link_quat_w`.""" return self.body_link_quat_w @property - def body_vel_w(self) -> wp.array: + def body_vel_w(self) -> ProxyArray: """Shorthand for :attr:`body_com_vel_w`.""" return self.body_com_vel_w @property - def body_lin_vel_w(self) -> wp.array: + def body_lin_vel_w(self) -> ProxyArray: """Shorthand for :attr:`body_com_lin_vel_w`.""" return self.body_com_lin_vel_w @property - def body_ang_vel_w(self) -> wp.array: + def body_ang_vel_w(self) -> ProxyArray: """Shorthand for :attr:`body_com_ang_vel_w`.""" return self.body_com_ang_vel_w @property - def body_acc_w(self) -> wp.array: + def body_acc_w(self) -> ProxyArray: """Shorthand for :attr:`body_com_acc_w`.""" return self.body_com_acc_w @property - def body_lin_acc_w(self) -> wp.array: + def body_lin_acc_w(self) -> ProxyArray: """Shorthand for :attr:`body_com_lin_acc_w`.""" return self.body_com_lin_acc_w @property - def body_ang_acc_w(self) -> wp.array: + def body_ang_acc_w(self) -> ProxyArray: """Shorthand for :attr:`body_com_ang_acc_w`.""" return self.body_com_ang_acc_w @property - def com_pos_b(self) -> wp.array: + def com_pos_b(self) -> ProxyArray: """Shorthand for :attr:`body_com_pos_b`.""" return self.body_com_pos_b @property - def com_quat_b(self) -> wp.array: + def com_quat_b(self) -> ProxyArray: """Shorthand for :attr:`body_com_quat_b`.""" return self.body_com_quat_b @@ -682,7 +684,7 @@ def com_quat_b(self) -> wp.array: """ @property - def default_mass(self) -> wp.array: + def default_mass(self) -> ProxyArray: """Deprecated property. Please use :attr:`body_mass` instead and manage the default mass manually.""" warnings.warn( "The `default_mass` property will be deprecated in a IsaacLab 4.0. Please use `body_mass` instead. " @@ -691,11 +693,11 @@ def default_mass(self) -> wp.array: stacklevel=2, ) if self._default_mass is None: - self._default_mass = wp.clone(self.body_mass, self.device) - return self._default_mass + self._default_mass = wp.clone(self.body_mass.warp, self.device) + return ProxyArray(self._default_mass) @property - def default_inertia(self) -> wp.array: + def default_inertia(self) -> ProxyArray: """Deprecated property. Please use :attr:`body_inertia` instead and manage the default inertia manually.""" warnings.warn( "The `default_inertia` property will be deprecated in a IsaacLab 4.0. Please use `body_inertia` instead. " @@ -704,5 +706,5 @@ def default_inertia(self) -> wp.array: stacklevel=2, ) if self._default_inertia is None: - self._default_inertia = wp.clone(self.body_inertia, self.device) - return self._default_inertia + self._default_inertia = wp.clone(self.body_inertia.warp, self.device) + return ProxyArray(self._default_inertia) diff --git a/source/isaaclab/isaaclab/assets/rigid_object_collection/base_rigid_object_collection_data.py b/source/isaaclab/isaaclab/assets/rigid_object_collection/base_rigid_object_collection_data.py index 1913bb4bef9b..9ab4e718c47e 100644 --- a/source/isaaclab/isaaclab/assets/rigid_object_collection/base_rigid_object_collection_data.py +++ b/source/isaaclab/isaaclab/assets/rigid_object_collection/base_rigid_object_collection_data.py @@ -8,6 +8,8 @@ import warp as wp +from isaaclab.utils.warp import ProxyArray + class BaseRigidObjectCollectionData(ABC): """Data container for a rigid object collection. @@ -63,7 +65,7 @@ def update(self, dt: float) -> None: @property @abstractmethod - def default_body_pose(self) -> wp.array: + def default_body_pose(self) -> ProxyArray: """Default body pose ``[pos, quat]`` in local environment frame. The position and quaternion are of the rigid body's actor frame. @@ -74,7 +76,7 @@ def default_body_pose(self) -> wp.array: @property @abstractmethod - def default_body_vel(self) -> wp.array: + def default_body_vel(self) -> ProxyArray: """Default body velocity ``[lin_vel, ang_vel]`` in local environment frame. The linear and angular velocities are of the rigid body's center of mass frame. @@ -85,7 +87,7 @@ def default_body_vel(self) -> wp.array: @property @abstractmethod - def default_body_state(self) -> wp.array: + def default_body_state(self) -> ProxyArray: """Deprecated, same as :attr:`default_body_pose` and :attr:`default_body_vel`.""" raise NotImplementedError() @@ -95,7 +97,7 @@ def default_body_state(self) -> wp.array: @property @abstractmethod - def body_link_pose_w(self) -> wp.array: + def body_link_pose_w(self) -> ProxyArray: """Body link pose ``[pos, quat]`` in simulation world frame. Shape is (num_instances, num_bodies), dtype = wp.transformf. In torch this resolves to @@ -108,7 +110,7 @@ def body_link_pose_w(self) -> wp.array: @property @abstractmethod - def body_link_vel_w(self) -> wp.array: + def body_link_vel_w(self) -> ProxyArray: """Body link velocity ``[lin_vel, ang_vel]`` in simulation world frame. Shape is (num_instances, num_bodies), dtype = wp.spatial_vectorf. In torch this resolves to @@ -121,7 +123,7 @@ def body_link_vel_w(self) -> wp.array: @property @abstractmethod - def body_com_pose_w(self) -> wp.array: + def body_com_pose_w(self) -> ProxyArray: """Body center of mass pose ``[pos, quat]`` in simulation world frame. Shape is (num_instances, num_bodies), dtype = wp.transformf. In torch this resolves to @@ -134,7 +136,7 @@ def body_com_pose_w(self) -> wp.array: @property @abstractmethod - def body_com_vel_w(self) -> wp.array: + def body_com_vel_w(self) -> ProxyArray: """Body center of mass velocity ``[lin_vel, ang_vel]`` in simulation world frame. Shape is (num_instances, num_bodies), dtype = wp.spatial_vectorf. In torch this resolves to @@ -147,25 +149,25 @@ def body_com_vel_w(self) -> wp.array: @property @abstractmethod - def body_state_w(self) -> wp.array: + def body_state_w(self) -> ProxyArray: """Deprecated, same as :attr:`body_link_pose_w` and :attr:`body_com_vel_w`.""" raise NotImplementedError() @property @abstractmethod - def body_link_state_w(self) -> wp.array: + def body_link_state_w(self) -> ProxyArray: """Deprecated, same as :attr:`body_link_pose_w` and :attr:`body_link_vel_w`.""" raise NotImplementedError() @property @abstractmethod - def body_com_state_w(self) -> wp.array: + def body_com_state_w(self) -> ProxyArray: """Deprecated, same as :attr:`body_com_pose_w` and :attr:`body_com_vel_w`.""" raise NotImplementedError() @property @abstractmethod - def body_com_acc_w(self) -> wp.array: + def body_com_acc_w(self) -> ProxyArray: """Acceleration of all bodies ``[lin_acc, ang_acc]`` in the simulation world frame. Shape is (num_instances, num_bodies), dtype = wp.spatial_vectorf. In torch this resolves to @@ -177,7 +179,7 @@ def body_com_acc_w(self) -> wp.array: @property @abstractmethod - def body_com_pose_b(self) -> wp.array: + def body_com_pose_b(self) -> ProxyArray: """Center of mass pose ``[pos, quat]`` of all bodies in their respective body's link frames. Shape is (num_instances, num_bodies), dtype = wp.transformf. In torch this resolves to @@ -190,7 +192,7 @@ def body_com_pose_b(self) -> wp.array: @property @abstractmethod - def body_mass(self) -> wp.array: + def body_mass(self) -> ProxyArray: """Mass of all bodies in the simulation world frame. Shape is (num_instances, num_bodies), dtype = wp.float32. In torch this resolves to (num_instances, num_bodies). @@ -199,7 +201,7 @@ def body_mass(self) -> wp.array: @property @abstractmethod - def body_inertia(self) -> wp.array: + def body_inertia(self) -> ProxyArray: """Inertia of all bodies in the simulation world frame. Shape is (num_instances, num_bodies, 9), dtype = wp.float32. In torch this resolves to @@ -213,7 +215,7 @@ def body_inertia(self) -> wp.array: @property @abstractmethod - def projected_gravity_b(self) -> wp.array: + def projected_gravity_b(self) -> ProxyArray: """Projection of the gravity direction on base frame. Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to @@ -223,7 +225,7 @@ def projected_gravity_b(self) -> wp.array: @property @abstractmethod - def heading_w(self) -> wp.array: + def heading_w(self) -> ProxyArray: """Yaw heading of the base frame (in radians). Shape is (num_instances, num_bodies), dtype = wp.float32. @@ -237,7 +239,7 @@ def heading_w(self) -> wp.array: @property @abstractmethod - def body_link_lin_vel_b(self) -> wp.array: + def body_link_lin_vel_b(self) -> ProxyArray: """Root link linear velocity in base frame. Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to @@ -250,7 +252,7 @@ def body_link_lin_vel_b(self) -> wp.array: @property @abstractmethod - def body_link_ang_vel_b(self) -> wp.array: + def body_link_ang_vel_b(self) -> ProxyArray: """Root link angular velocity in base frame. Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to @@ -263,7 +265,7 @@ def body_link_ang_vel_b(self) -> wp.array: @property @abstractmethod - def body_com_lin_vel_b(self) -> wp.array: + def body_com_lin_vel_b(self) -> ProxyArray: """Root center of mass linear velocity in base frame. Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to @@ -276,7 +278,7 @@ def body_com_lin_vel_b(self) -> wp.array: @property @abstractmethod - def body_com_ang_vel_b(self) -> wp.array: + def body_com_ang_vel_b(self) -> ProxyArray: """Root center of mass angular velocity in base frame. Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to @@ -293,7 +295,7 @@ def body_com_ang_vel_b(self) -> wp.array: @property @abstractmethod - def body_link_pos_w(self) -> wp.array: + def body_link_pos_w(self) -> ProxyArray: """Positions of all bodies in simulation world frame. Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to @@ -305,7 +307,7 @@ def body_link_pos_w(self) -> wp.array: @property @abstractmethod - def body_link_quat_w(self) -> wp.array: + def body_link_quat_w(self) -> ProxyArray: """Orientation (x, y, z, w) of all bodies in simulation world frame. Shape is (num_instances, num_bodies), dtype = wp.quatf. In torch this resolves to @@ -317,7 +319,7 @@ def body_link_quat_w(self) -> wp.array: @property @abstractmethod - def body_link_lin_vel_w(self) -> wp.array: + def body_link_lin_vel_w(self) -> ProxyArray: """Linear velocity of all bodies in simulation world frame. Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to @@ -329,7 +331,7 @@ def body_link_lin_vel_w(self) -> wp.array: @property @abstractmethod - def body_link_ang_vel_w(self) -> wp.array: + def body_link_ang_vel_w(self) -> ProxyArray: """Angular velocity of all bodies in simulation world frame. Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to @@ -341,7 +343,7 @@ def body_link_ang_vel_w(self) -> wp.array: @property @abstractmethod - def body_com_pos_w(self) -> wp.array: + def body_com_pos_w(self) -> ProxyArray: """Positions of all bodies' center of mass in simulation world frame. Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to @@ -353,7 +355,7 @@ def body_com_pos_w(self) -> wp.array: @property @abstractmethod - def body_com_quat_w(self) -> wp.array: + def body_com_quat_w(self) -> ProxyArray: """Orientation (x, y, z, w) of the principal axes of inertia of all bodies in simulation world frame. Shape is (num_instances, num_bodies), dtype = wp.quatf. In torch this resolves to @@ -365,7 +367,7 @@ def body_com_quat_w(self) -> wp.array: @property @abstractmethod - def body_com_lin_vel_w(self) -> wp.array: + def body_com_lin_vel_w(self) -> ProxyArray: """Linear velocity of all bodies in simulation world frame. Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to @@ -377,7 +379,7 @@ def body_com_lin_vel_w(self) -> wp.array: @property @abstractmethod - def body_com_ang_vel_w(self) -> wp.array: + def body_com_ang_vel_w(self) -> ProxyArray: """Angular velocity of all bodies in simulation world frame. Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to @@ -389,7 +391,7 @@ def body_com_ang_vel_w(self) -> wp.array: @property @abstractmethod - def body_com_lin_acc_w(self) -> wp.array: + def body_com_lin_acc_w(self) -> ProxyArray: """Linear acceleration of all bodies in simulation world frame. Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to @@ -401,7 +403,7 @@ def body_com_lin_acc_w(self) -> wp.array: @property @abstractmethod - def body_com_ang_acc_w(self) -> wp.array: + def body_com_ang_acc_w(self) -> ProxyArray: """Angular acceleration of all bodies in simulation world frame. Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to @@ -413,7 +415,7 @@ def body_com_ang_acc_w(self) -> wp.array: @property @abstractmethod - def body_com_pos_b(self) -> wp.array: + def body_com_pos_b(self) -> ProxyArray: """Center of mass position of all of the bodies in their respective link frames. Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to @@ -425,7 +427,7 @@ def body_com_pos_b(self) -> wp.array: @property @abstractmethod - def body_com_quat_b(self) -> wp.array: + def body_com_quat_b(self) -> ProxyArray: """Orientation (x, y, z, w) of the principal axes of inertia of all of the bodies in their respective link frames. @@ -441,57 +443,57 @@ def body_com_quat_b(self) -> wp.array: """ @property - def body_pose_w(self) -> wp.array: + def body_pose_w(self) -> ProxyArray: """Shorthand for :attr:`body_link_pose_w`.""" return self.body_link_pose_w @property - def body_pos_w(self) -> wp.array: + def body_pos_w(self) -> ProxyArray: """Shorthand for :attr:`body_link_pos_w`.""" return self.body_link_pos_w @property - def body_quat_w(self) -> wp.array: + def body_quat_w(self) -> ProxyArray: """Shorthand for :attr:`body_link_quat_w`.""" return self.body_link_quat_w @property - def body_vel_w(self) -> wp.array: + def body_vel_w(self) -> ProxyArray: """Shorthand for :attr:`body_com_vel_w`.""" return self.body_com_vel_w @property - def body_lin_vel_w(self) -> wp.array: + def body_lin_vel_w(self) -> ProxyArray: """Shorthand for :attr:`body_com_lin_vel_w`.""" return self.body_com_lin_vel_w @property - def body_ang_vel_w(self) -> wp.array: + def body_ang_vel_w(self) -> ProxyArray: """Shorthand for :attr:`body_com_ang_vel_w`.""" return self.body_com_ang_vel_w @property - def body_acc_w(self) -> wp.array: + def body_acc_w(self) -> ProxyArray: """Shorthand for :attr:`body_com_acc_w`.""" return self.body_com_acc_w @property - def body_lin_acc_w(self) -> wp.array: + def body_lin_acc_w(self) -> ProxyArray: """Shorthand for :attr:`body_com_lin_acc_w`.""" return self.body_com_lin_acc_w @property - def body_ang_acc_w(self) -> wp.array: + def body_ang_acc_w(self) -> ProxyArray: """Shorthand for :attr:`body_com_ang_acc_w`.""" return self.body_com_ang_acc_w @property - def com_pos_b(self) -> wp.array: + def com_pos_b(self) -> ProxyArray: """Shorthand for :attr:`body_com_pos_b`.""" return self.body_com_pos_b @property - def com_quat_b(self) -> wp.array: + def com_quat_b(self) -> ProxyArray: """Shorthand for :attr:`body_com_quat_b`.""" return self.body_com_quat_b @@ -505,7 +507,7 @@ def _create_buffers(self): """ @property - def default_object_pose(self) -> wp.array: + def default_object_pose(self) -> ProxyArray: """Deprecated property. Please use :attr:`default_body_pose` instead.""" warnings.warn( "The `default_object_pose` property will be deprecated in a IsaacLab 4.0. Please use" @@ -516,7 +518,7 @@ def default_object_pose(self) -> wp.array: return self.default_body_pose @property - def default_object_vel(self) -> wp.array: + def default_object_vel(self) -> ProxyArray: """Deprecated property. Please use :attr:`default_body_vel` instead.""" warnings.warn( "The `default_object_vel` property will be deprecated in a IsaacLab 4.0. Please use" @@ -527,7 +529,7 @@ def default_object_vel(self) -> wp.array: return self.default_body_vel @property - def default_object_state(self) -> wp.array: + def default_object_state(self) -> ProxyArray: """Deprecated property. Please use :attr:`default_body_state` instead.""" warnings.warn( "The `default_object_state` property will be deprecated in a IsaacLab 4.0. Please use" @@ -636,7 +638,7 @@ def object_com_pose_b(self): return self.body_com_pose_b @property - def object_link_pos_w(self) -> wp.array: + def object_link_pos_w(self) -> ProxyArray: """Deprecated property. Please use :attr:`body_link_pos_w` instead.""" warnings.warn( "The `object_link_pos_w` property will be deprecated in a IsaacLab 4.0. Please use" @@ -647,7 +649,7 @@ def object_link_pos_w(self) -> wp.array: return self.body_link_pos_w @property - def object_link_quat_w(self) -> wp.array: + def object_link_quat_w(self) -> ProxyArray: """Deprecated property. Please use :attr:`body_link_quat_w` instead.""" warnings.warn( "The `object_link_quat_w` property will be deprecated in a IsaacLab 4.0. Please use" @@ -658,7 +660,7 @@ def object_link_quat_w(self) -> wp.array: return self.body_link_quat_w @property - def object_link_lin_vel_w(self) -> wp.array: + def object_link_lin_vel_w(self) -> ProxyArray: """Deprecated property. Please use :attr:`body_link_lin_vel_w` instead.""" warnings.warn( "The `object_link_lin_vel_w` property will be deprecated in a IsaacLab 4.0. Please use" @@ -669,7 +671,7 @@ def object_link_lin_vel_w(self) -> wp.array: return self.body_link_lin_vel_w @property - def object_link_ang_vel_w(self) -> wp.array: + def object_link_ang_vel_w(self) -> ProxyArray: """Deprecated property. Please use :attr:`body_link_ang_vel_w` instead.""" warnings.warn( "The `object_link_ang_vel_w` property will be deprecated in a IsaacLab 4.0. Please use" @@ -680,7 +682,7 @@ def object_link_ang_vel_w(self) -> wp.array: return self.body_link_ang_vel_w @property - def object_com_pos_w(self) -> wp.array: + def object_com_pos_w(self) -> ProxyArray: """Deprecated property. Please use :attr:`body_com_pos_w` instead.""" warnings.warn( "The `object_com_pos_w` property will be deprecated in a IsaacLab 4.0. Please use" @@ -691,7 +693,7 @@ def object_com_pos_w(self) -> wp.array: return self.body_com_pos_w @property - def object_com_quat_w(self) -> wp.array: + def object_com_quat_w(self) -> ProxyArray: """Deprecated property. Please use :attr:`body_com_quat_w` instead.""" warnings.warn( "The `object_com_quat_w` property will be deprecated in a IsaacLab 4.0. Please use" @@ -702,7 +704,7 @@ def object_com_quat_w(self) -> wp.array: return self.body_com_quat_w @property - def object_com_lin_vel_w(self) -> wp.array: + def object_com_lin_vel_w(self) -> ProxyArray: """Deprecated property. Please use :attr:`body_com_lin_vel_w` instead.""" warnings.warn( "The `object_com_lin_vel_w` property will be deprecated in a IsaacLab 4.0. Please use" @@ -713,7 +715,7 @@ def object_com_lin_vel_w(self) -> wp.array: return self.body_com_lin_vel_w @property - def object_com_ang_vel_w(self) -> wp.array: + def object_com_ang_vel_w(self) -> ProxyArray: """Deprecated property. Please use :attr:`body_com_ang_vel_w` instead.""" warnings.warn( "The `object_com_ang_vel_w` property will be deprecated in a IsaacLab 4.0. Please use" @@ -724,7 +726,7 @@ def object_com_ang_vel_w(self) -> wp.array: return self.body_com_ang_vel_w @property - def object_com_lin_acc_w(self) -> wp.array: + def object_com_lin_acc_w(self) -> ProxyArray: """Deprecated property. Please use :attr:`body_com_lin_acc_w` instead.""" warnings.warn( "The `object_com_lin_acc_w` property will be deprecated in a IsaacLab 4.0. Please use" @@ -735,7 +737,7 @@ def object_com_lin_acc_w(self) -> wp.array: return self.body_com_lin_acc_w @property - def object_com_ang_acc_w(self) -> wp.array: + def object_com_ang_acc_w(self) -> ProxyArray: """Deprecated property. Please use :attr:`body_com_ang_acc_w` instead.""" warnings.warn( "The `object_com_ang_acc_w` property will be deprecated in a IsaacLab 4.0. Please use" @@ -746,7 +748,7 @@ def object_com_ang_acc_w(self) -> wp.array: return self.body_com_ang_acc_w @property - def object_com_pos_b(self) -> wp.array: + def object_com_pos_b(self) -> ProxyArray: """Deprecated property. Please use :attr:`body_com_pos_b` instead.""" warnings.warn( "The `object_com_pos_b` property will be deprecated in a IsaacLab 4.0. Please use" @@ -757,7 +759,7 @@ def object_com_pos_b(self) -> wp.array: return self.body_com_pos_b @property - def object_com_quat_b(self) -> wp.array: + def object_com_quat_b(self) -> ProxyArray: """Deprecated property. Please use :attr:`body_com_quat_b` instead.""" warnings.warn( "The `object_com_quat_b` property will be deprecated in a IsaacLab 4.0. Please use" @@ -768,7 +770,7 @@ def object_com_quat_b(self) -> wp.array: return self.body_com_quat_b @property - def object_link_lin_vel_b(self) -> wp.array: + def object_link_lin_vel_b(self) -> ProxyArray: """Deprecated property. Please use :attr:`body_link_lin_vel_b` instead.""" warnings.warn( "The `object_link_lin_vel_b` property will be deprecated in a IsaacLab 4.0. Please use" @@ -779,7 +781,7 @@ def object_link_lin_vel_b(self) -> wp.array: return self.body_link_lin_vel_b @property - def object_link_ang_vel_b(self) -> wp.array: + def object_link_ang_vel_b(self) -> ProxyArray: """Deprecated property. Please use :attr:`body_link_ang_vel_b` instead.""" warnings.warn( "The `object_link_ang_vel_b` property will be deprecated in a IsaacLab 4.0. Please use" @@ -790,7 +792,7 @@ def object_link_ang_vel_b(self) -> wp.array: return self.body_link_ang_vel_b @property - def object_com_lin_vel_b(self) -> wp.array: + def object_com_lin_vel_b(self) -> ProxyArray: """Deprecated property. Please use :attr:`body_com_lin_vel_b` instead.""" warnings.warn( "The `object_com_lin_vel_b` property will be deprecated in a IsaacLab 4.0. Please use" @@ -801,7 +803,7 @@ def object_com_lin_vel_b(self) -> wp.array: return self.body_com_lin_vel_b @property - def object_com_ang_vel_b(self) -> wp.array: + def object_com_ang_vel_b(self) -> ProxyArray: """Deprecated property. Please use :attr:`body_com_ang_vel_b` instead.""" warnings.warn( "The `object_com_ang_vel_b` property will be deprecated in a IsaacLab 4.0. Please use" @@ -812,7 +814,7 @@ def object_com_ang_vel_b(self) -> wp.array: return self.body_com_ang_vel_b @property - def object_pose_w(self) -> wp.array: + def object_pose_w(self) -> ProxyArray: """Deprecated property. Please use :attr:`body_link_pose_w` instead.""" warnings.warn( "The `object_pose_w` property will be deprecated in a IsaacLab 4.0. Please use `body_link_pose_w` instead.", @@ -822,7 +824,7 @@ def object_pose_w(self) -> wp.array: return self.body_link_pose_w @property - def object_pos_w(self) -> wp.array: + def object_pos_w(self) -> ProxyArray: """Deprecated property. Please use :attr:`body_link_pos_w` instead.""" warnings.warn( "The `object_pos_w` property will be deprecated in a IsaacLab 4.0. Please use `body_link_pos_w` instead.", @@ -832,7 +834,7 @@ def object_pos_w(self) -> wp.array: return self.body_link_pos_w @property - def object_quat_w(self) -> wp.array: + def object_quat_w(self) -> ProxyArray: """Deprecated property. Please use :attr:`body_link_quat_w` instead.""" warnings.warn( "The `object_quat_w` property will be deprecated in a IsaacLab 4.0. Please use `body_link_quat_w` instead.", @@ -842,7 +844,7 @@ def object_quat_w(self) -> wp.array: return self.body_link_quat_w @property - def object_vel_w(self) -> wp.array: + def object_vel_w(self) -> ProxyArray: """Deprecated property. Please use :attr:`body_com_vel_w` instead.""" warnings.warn( "The `object_vel_w` property will be deprecated in a IsaacLab 4.0. Please use `body_com_vel_w` instead.", @@ -852,7 +854,7 @@ def object_vel_w(self) -> wp.array: return self.body_com_vel_w @property - def object_lin_vel_w(self) -> wp.array: + def object_lin_vel_w(self) -> ProxyArray: """Deprecated property. Please use :attr:`body_com_lin_vel_w` instead.""" warnings.warn( "The `object_lin_vel_w` property will be deprecated in a IsaacLab 4.0. Please use" @@ -863,7 +865,7 @@ def object_lin_vel_w(self) -> wp.array: return self.body_com_lin_vel_w @property - def object_ang_vel_w(self) -> wp.array: + def object_ang_vel_w(self) -> ProxyArray: """Deprecated property. Please use :attr:`body_com_ang_vel_w` instead.""" warnings.warn( "The `object_ang_vel_w` property will be deprecated in a IsaacLab 4.0. Please use" @@ -874,7 +876,7 @@ def object_ang_vel_w(self) -> wp.array: return self.body_com_ang_vel_w @property - def object_lin_vel_b(self) -> wp.array: + def object_lin_vel_b(self) -> ProxyArray: """Deprecated property. Please use :attr:`body_com_lin_vel_b` instead.""" warnings.warn( "The `object_lin_vel_b` property will be deprecated in a IsaacLab 4.0. Please use" @@ -885,7 +887,7 @@ def object_lin_vel_b(self) -> wp.array: return self.body_com_lin_vel_b @property - def object_ang_vel_b(self) -> wp.array: + def object_ang_vel_b(self) -> ProxyArray: """Deprecated property. Please use :attr:`body_com_ang_vel_b` instead.""" warnings.warn( "The `object_ang_vel_b` property will be deprecated in a IsaacLab 4.0. Please use" @@ -896,7 +898,7 @@ def object_ang_vel_b(self) -> wp.array: return self.body_com_ang_vel_b @property - def object_acc_w(self) -> wp.array: + def object_acc_w(self) -> ProxyArray: """Deprecated property. Please use :attr:`body_com_acc_w` instead.""" warnings.warn( "The `object_acc_w` property will be deprecated in a IsaacLab 4.0. Please use `body_com_acc_w` instead.", @@ -906,7 +908,7 @@ def object_acc_w(self) -> wp.array: return self.body_com_acc_w @property - def object_lin_acc_w(self) -> wp.array: + def object_lin_acc_w(self) -> ProxyArray: """Deprecated property. Please use :attr:`body_com_lin_acc_w` instead.""" warnings.warn( "The `object_lin_acc_w` property will be deprecated in a IsaacLab 4.0. Please use" @@ -917,7 +919,7 @@ def object_lin_acc_w(self) -> wp.array: return self.body_com_lin_acc_w @property - def object_ang_acc_w(self) -> wp.array: + def object_ang_acc_w(self) -> ProxyArray: """Deprecated property. Please use :attr:`body_com_ang_acc_w` instead.""" warnings.warn( "The `object_ang_acc_w` property will be deprecated in a IsaacLab 4.0. Please use" @@ -932,7 +934,7 @@ def object_ang_acc_w(self) -> wp.array: """ @property - def default_mass(self) -> wp.array: + def default_mass(self) -> ProxyArray: """Deprecated property. Please use :attr:`body_mass` instead and manage the default mass manually.""" warnings.warn( "The `default_mass` property will be deprecated in a IsaacLab 4.0. Please use `body_mass` instead. " @@ -941,11 +943,11 @@ def default_mass(self) -> wp.array: stacklevel=2, ) if self._default_mass is None: - self._default_mass = wp.clone(self.body_mass, self.device) - return self._default_mass + self._default_mass = wp.clone(self.body_mass.warp, self.device) + return ProxyArray(self._default_mass) @property - def default_inertia(self) -> wp.array: + def default_inertia(self) -> ProxyArray: """Deprecated property. Please use :attr:`body_inertia` instead and manage the default inertia manually.""" warnings.warn( "The `default_inertia` property will be deprecated in a IsaacLab 4.0. Please use `body_inertia` instead. " @@ -954,5 +956,5 @@ def default_inertia(self) -> wp.array: stacklevel=2, ) if self._default_inertia is None: - self._default_inertia = wp.clone(self.body_inertia, self.device) - return self._default_inertia + self._default_inertia = wp.clone(self.body_inertia.warp, self.device) + return ProxyArray(self._default_inertia) diff --git a/source/isaaclab/isaaclab/envs/mdp/actions/joint_actions.py b/source/isaaclab/isaaclab/envs/mdp/actions/joint_actions.py index 8fbeb996bfa5..ca0041213c78 100644 --- a/source/isaaclab/isaaclab/envs/mdp/actions/joint_actions.py +++ b/source/isaaclab/isaaclab/envs/mdp/actions/joint_actions.py @@ -10,7 +10,6 @@ from typing import TYPE_CHECKING import torch -import warp as wp import isaaclab.utils.string as string_utils from isaaclab.assets.articulation import Articulation @@ -193,7 +192,7 @@ def __init__(self, cfg: actions_cfg.JointPositionActionCfg, env: ManagerBasedEnv super().__init__(cfg, env) # use default joint positions as offset if cfg.use_default_offset: - self._offset = wp.to_torch(self._asset.data.default_joint_pos)[:, self._joint_ids].clone() + self._offset = self._asset.data.default_joint_pos.torch[:, self._joint_ids].clone() def apply_actions(self): # set position targets @@ -228,7 +227,7 @@ def __init__(self, cfg: actions_cfg.RelativeJointPositionActionCfg, env: Manager def apply_actions(self): # add current joint positions to the processed actions - current_actions = self.processed_actions + wp.to_torch(self._asset.data.joint_pos)[:, self._joint_ids] + current_actions = self.processed_actions + self._asset.data.joint_pos.torch[:, self._joint_ids] # set position targets self._asset.set_joint_position_target_index(target=current_actions, joint_ids=self._joint_ids) @@ -244,7 +243,7 @@ def __init__(self, cfg: actions_cfg.JointVelocityActionCfg, env: ManagerBasedEnv super().__init__(cfg, env) # use default joint velocity as offset if cfg.use_default_offset: - self._offset = wp.to_torch(self._asset.data.default_joint_vel)[:, self._joint_ids].clone() + self._offset = self._asset.data.default_joint_vel.torch[:, self._joint_ids].clone() def apply_actions(self): # set joint velocity targets diff --git a/source/isaaclab/isaaclab/envs/mdp/actions/joint_actions_to_limits.py b/source/isaaclab/isaaclab/envs/mdp/actions/joint_actions_to_limits.py index 488975e08187..7b77700cc6e3 100644 --- a/source/isaaclab/isaaclab/envs/mdp/actions/joint_actions_to_limits.py +++ b/source/isaaclab/isaaclab/envs/mdp/actions/joint_actions_to_limits.py @@ -10,7 +10,6 @@ from typing import TYPE_CHECKING import torch -import warp as wp import isaaclab.utils.math as math_utils import isaaclab.utils.string as string_utils @@ -169,8 +168,8 @@ def process_actions(self, actions: torch.Tensor): # rescale within the joint limits actions = math_utils.unscale_transform( actions, - wp.to_torch(self._asset.data.soft_joint_pos_limits)[:, self._joint_ids, 0], - wp.to_torch(self._asset.data.soft_joint_pos_limits)[:, self._joint_ids, 1], + self._asset.data.soft_joint_pos_limits.torch[:, self._joint_ids, 0], + self._asset.data.soft_joint_pos_limits.torch[:, self._joint_ids, 1], ) self._processed_actions[:] = actions[:] @@ -271,10 +270,10 @@ def reset(self, env_ids: Sequence[int] | None = None) -> None: # check if specific environment ids are provided if env_ids is None: super().reset(slice(None)) - self._prev_applied_actions[:] = wp.to_torch(self._asset.data.joint_pos)[:, self._joint_ids] + self._prev_applied_actions[:] = self._asset.data.joint_pos.torch[:, self._joint_ids] else: super().reset(env_ids) - curr_applied_actions = wp.to_torch(self._asset.data.joint_pos)[env_ids[:, None], self._joint_ids].view( + curr_applied_actions = self._asset.data.joint_pos.torch[env_ids[:, None], self._joint_ids].view( len(env_ids), -1 ) self._prev_applied_actions[env_ids, :] = curr_applied_actions @@ -288,8 +287,8 @@ def process_actions(self, actions: torch.Tensor): # clamp the targets self._processed_actions[:] = torch.clamp( ema_actions, - wp.to_torch(self._asset.data.soft_joint_pos_limits)[:, self._joint_ids, 0], - wp.to_torch(self._asset.data.soft_joint_pos_limits)[:, self._joint_ids, 1], + self._asset.data.soft_joint_pos_limits.torch[:, self._joint_ids, 0], + self._asset.data.soft_joint_pos_limits.torch[:, self._joint_ids, 1], ) # update previous targets self._prev_applied_actions[:] = self._processed_actions[:] diff --git a/source/isaaclab/isaaclab/envs/mdp/actions/non_holonomic_actions.py b/source/isaaclab/isaaclab/envs/mdp/actions/non_holonomic_actions.py index b4c88e3c171b..08d5e9be1c85 100644 --- a/source/isaaclab/isaaclab/envs/mdp/actions/non_holonomic_actions.py +++ b/source/isaaclab/isaaclab/envs/mdp/actions/non_holonomic_actions.py @@ -184,7 +184,7 @@ def process_actions(self, actions): def apply_actions(self): # obtain current heading - quat_w = self._asset.data.body_quat_w[:, self._body_idx].view(self.num_envs, 4) + quat_w = self._asset.data.body_quat_w.torch[:, self._body_idx].view(self.num_envs, 4) yaw_w = euler_xyz_from_quat(quat_w)[2] # compute joint velocities targets self._joint_vel_command[:, 0] = torch.cos(yaw_w) * self.processed_actions[:, 0] # x diff --git a/source/isaaclab/isaaclab/envs/mdp/actions/pink_task_space_actions.py b/source/isaaclab/isaaclab/envs/mdp/actions/pink_task_space_actions.py index 5d7e59418eea..3cb4450805b2 100644 --- a/source/isaaclab/isaaclab/envs/mdp/actions/pink_task_space_actions.py +++ b/source/isaaclab/isaaclab/envs/mdp/actions/pink_task_space_actions.py @@ -218,7 +218,7 @@ def _get_base_link_frame_transform(self) -> torch.Tensor: """ # Get base link frame pose in world origin using cached index articulation_data = self._env.scene[self.cfg.controller.articulation_name].data - base_link_frame_in_world_origin = wp.to_torch(articulation_data.body_link_state_w)[:, self._base_link_idx, :7] + base_link_frame_in_world_origin = articulation_data.body_link_state_w.torch[:, self._base_link_idx, :7] # Transform to environment origin frame (reuse buffer to avoid allocation) torch.sub( @@ -352,7 +352,7 @@ def _compute_ik_solutions(self) -> torch.Tensor: for env_index, ik_controller in enumerate(self._ik_controllers): # Get current joint positions for this environment - current_joint_pos = wp.to_torch(self._asset.data.joint_pos).cpu().numpy()[env_index] + current_joint_pos = self._asset.data.joint_pos.torch.cpu().numpy()[env_index] # Compute IK solution joint_pos_des = ik_controller.compute(current_joint_pos, self._sim_dt) diff --git a/source/isaaclab/isaaclab/envs/mdp/actions/rmpflow_task_space_actions.py b/source/isaaclab/isaaclab/envs/mdp/actions/rmpflow_task_space_actions.py index 741de96808c9..922db073a300 100644 --- a/source/isaaclab/isaaclab/envs/mdp/actions/rmpflow_task_space_actions.py +++ b/source/isaaclab/isaaclab/envs/mdp/actions/rmpflow_task_space_actions.py @@ -133,7 +133,7 @@ def jacobian_w(self) -> torch.Tensor: @property def jacobian_b(self) -> torch.Tensor: jacobian = self.jacobian_w - base_rot = wp.to_torch(self._asset.data.root_quat_w) + base_rot = self._asset.data.root_quat_w.torch base_rot_matrix = math_utils.matrix_from_quat(math_utils.quat_inv(base_rot)) jacobian[:, :3, :] = torch.bmm(base_rot_matrix, jacobian[:, :3, :]) jacobian[:, 3:, :] = torch.bmm(base_rot_matrix, jacobian[:, 3:, :]) @@ -178,11 +178,11 @@ def process_actions(self, actions: torch.Tensor): def apply_actions(self): # obtain quantities from simulation ee_pos_curr, ee_quat_curr = self._compute_frame_pose() - joint_pos = wp.to_torch(self._asset.data.joint_pos)[:, self._joint_ids] + joint_pos = self._asset.data.joint_pos.torch[:, self._joint_ids] # compute the delta in joint-space if ee_quat_curr.norm() != 0: joint_pos_des, joint_vel_des = self._rmpflow_controller.compute( - wp.to_torch(self._asset.data.joint_pos), wp.to_torch(self._asset.data.joint_vel) + self._asset.data.joint_pos.torch, self._asset.data.joint_vel.torch ) else: joint_pos_des = joint_pos.clone() @@ -206,10 +206,10 @@ def _compute_frame_pose(self) -> tuple[torch.Tensor, torch.Tensor]: A tuple of the body's position and orientation in the root frame. """ # obtain quantities from simulation - ee_pos_w = wp.to_torch(self._asset.data.body_pos_w)[:, self._body_idx] - ee_quat_w = wp.to_torch(self._asset.data.body_quat_w)[:, self._body_idx] - root_pos_w = wp.to_torch(self._asset.data.root_pos_w) - root_quat_w = wp.to_torch(self._asset.data.root_quat_w) + ee_pos_w = self._asset.data.body_pos_w.torch[:, self._body_idx] + ee_quat_w = self._asset.data.body_quat_w.torch[:, self._body_idx] + root_pos_w = self._asset.data.root_pos_w.torch + root_quat_w = self._asset.data.root_quat_w.torch # compute the pose of the body in the root frame ee_pose_b, ee_quat_b = math_utils.subtract_frame_transforms(root_pos_w, root_quat_w, ee_pos_w, ee_quat_w) # account for the offset diff --git a/source/isaaclab/isaaclab/envs/mdp/actions/task_space_actions.py b/source/isaaclab/isaaclab/envs/mdp/actions/task_space_actions.py index 99ee2e5cfb55..4aa66367d2c2 100644 --- a/source/isaaclab/isaaclab/envs/mdp/actions/task_space_actions.py +++ b/source/isaaclab/isaaclab/envs/mdp/actions/task_space_actions.py @@ -148,7 +148,7 @@ def jacobian_w(self) -> torch.Tensor: @property def jacobian_b(self) -> torch.Tensor: jacobian = self.jacobian_w - base_rot = wp.to_torch(self._asset.data.root_quat_w) + base_rot = self._asset.data.root_quat_w.torch base_rot_matrix = math_utils.matrix_from_quat(math_utils.quat_inv(base_rot)) jacobian[:, :3, :] = torch.bmm(base_rot_matrix, jacobian[:, :3, :]) jacobian[:, 3:, :] = torch.bmm(base_rot_matrix, jacobian[:, 3:, :]) @@ -205,7 +205,7 @@ def process_actions(self, actions: torch.Tensor): def apply_actions(self): # obtain quantities from simulation ee_pos_curr, ee_quat_curr = self._compute_frame_pose() - joint_pos = wp.to_torch(self._asset.data.joint_pos)[:, self._joint_ids] + joint_pos = self._asset.data.joint_pos.torch[:, self._joint_ids] # compute the delta in joint-space if ee_quat_curr.norm() != 0: jacobian = self._compute_frame_jacobian() @@ -229,10 +229,10 @@ def _compute_frame_pose(self) -> tuple[torch.Tensor, torch.Tensor]: A tuple of the body's position and orientation in the root frame. """ # obtain quantities from simulation - ee_pos_w = wp.to_torch(self._asset.data.body_pos_w)[:, self._body_idx] - ee_quat_w = wp.to_torch(self._asset.data.body_quat_w)[:, self._body_idx] - root_pos_w = wp.to_torch(self._asset.data.root_pos_w) - root_quat_w = wp.to_torch(self._asset.data.root_quat_w) + ee_pos_w = self._asset.data.body_pos_w.torch[:, self._body_idx] + ee_quat_w = self._asset.data.body_quat_w.torch[:, self._body_idx] + root_pos_w = self._asset.data.root_pos_w.torch + root_quat_w = self._asset.data.root_quat_w.torch # compute the pose of the body in the root frame ee_pose_b, ee_quat_b = math_utils.subtract_frame_transforms(root_pos_w, root_quat_w, ee_pos_w, ee_quat_w) # account for the offset @@ -442,7 +442,7 @@ def jacobian_w(self) -> torch.Tensor: @property def jacobian_b(self) -> torch.Tensor: jacobian = self.jacobian_w - base_rot = wp.to_torch(self._asset.data.root_quat_w) + base_rot = self._asset.data.root_quat_w.torch base_rot_matrix = math_utils.matrix_from_quat(math_utils.quat_inv(base_rot)) jacobian[:, :3, :] = torch.bmm(base_rot_matrix, jacobian[:, :3, :]) jacobian[:, 3:, :] = torch.bmm(base_rot_matrix, jacobian[:, 3:, :]) @@ -637,11 +637,11 @@ def _resolve_nullspace_joint_pos_targets(self): elif self.cfg.nullspace_joint_pos_target == "center": # Get the center of the robot soft joint limits self._nullspace_joint_pos_target = torch.mean( - wp.to_torch(self._asset.data.soft_joint_pos_limits)[:, self._joint_ids, :], dim=-1 + self._asset.data.soft_joint_pos_limits.torch[:, self._joint_ids, :], dim=-1 ) elif self.cfg.nullspace_joint_pos_target == "default": # Get the default joint positions - self._nullspace_joint_pos_target = wp.to_torch(self._asset.data.default_joint_pos)[:, self._joint_ids] + self._nullspace_joint_pos_target = self._asset.data.default_joint_pos.torch[:, self._joint_ids] else: raise ValueError("Invalid value for nullspace joint pos targets.") @@ -689,12 +689,12 @@ def _compute_ee_jacobian(self): def _compute_ee_pose(self): """Computes the pose of the ee frame in root frame.""" # Obtain quantities from simulation - self._ee_pose_w[:, 0:3] = wp.to_torch(self._asset.data.body_pos_w)[:, self._ee_body_idx] - self._ee_pose_w[:, 3:7] = wp.to_torch(self._asset.data.body_quat_w)[:, self._ee_body_idx] + self._ee_pose_w[:, 0:3] = self._asset.data.body_pos_w.torch[:, self._ee_body_idx] + self._ee_pose_w[:, 3:7] = self._asset.data.body_quat_w.torch[:, self._ee_body_idx] # Compute the pose of the ee body in the root frame self._ee_pose_b_no_offset[:, 0:3], self._ee_pose_b_no_offset[:, 3:7] = math_utils.subtract_frame_transforms( - wp.to_torch(self._asset.data.root_pos_w), - wp.to_torch(self._asset.data.root_quat_w), + self._asset.data.root_pos_w.torch, + self._asset.data.root_quat_w.torch, self._ee_pose_w[:, 0:3], self._ee_pose_w[:, 3:7], ) @@ -709,12 +709,12 @@ def _compute_ee_pose(self): def _compute_ee_velocity(self): """Computes the velocity of the ee frame in root frame.""" # Extract end-effector velocity in the world frame - self._ee_vel_w[:] = wp.to_torch(self._asset.data.body_vel_w)[:, self._ee_body_idx, :] + self._ee_vel_w[:] = self._asset.data.body_vel_w.torch[:, self._ee_body_idx, :] # Compute the relative velocity in the world frame - relative_vel_w = self._ee_vel_w - wp.to_torch(self._asset.data.root_vel_w) + relative_vel_w = self._ee_vel_w - self._asset.data.root_vel_w.torch # Convert ee velocities from world to root frame - root_quat_w = wp.to_torch(self._asset.data.root_quat_w) + root_quat_w = self._asset.data.root_quat_w.torch self._ee_vel_b[:, 0:3] = math_utils.quat_apply_inverse(root_quat_w, relative_vel_w[:, 0:3]) self._ee_vel_b[:, 3:6] = math_utils.quat_apply_inverse(root_quat_w, relative_vel_w[:, 3:6]) @@ -731,17 +731,15 @@ def _compute_ee_force(self): # Obtain contact forces only if the contact sensor is available if self._contact_sensor is not None: self._contact_sensor.update(self._sim_dt) - self._ee_force_w[:] = wp.to_torch(self._contact_sensor.data.net_forces_w)[:, 0, :] # type: ignore + self._ee_force_w[:] = self._contact_sensor.data.net_forces_w.torch[:, 0, :] # type: ignore # Rotate forces and torques into root frame - self._ee_force_b[:] = math_utils.quat_apply_inverse( - wp.to_torch(self._asset.data.root_quat_w), self._ee_force_w - ) + self._ee_force_b[:] = math_utils.quat_apply_inverse(self._asset.data.root_quat_w.torch, self._ee_force_w) def _compute_joint_states(self): """Computes the joint states for operational space control.""" # Extract joint positions and velocities - self._joint_pos[:] = wp.to_torch(self._asset.data.joint_pos)[:, self._joint_ids] - self._joint_vel[:] = wp.to_torch(self._asset.data.joint_vel)[:, self._joint_ids] + self._joint_pos[:] = self._asset.data.joint_pos.torch[:, self._joint_ids] + self._joint_vel[:] = self._asset.data.joint_vel.torch[:, self._joint_ids] def _compute_task_frame_pose(self): """Computes the pose of the task frame in root frame.""" @@ -750,10 +748,10 @@ def _compute_task_frame_pose(self): self._task_frame_transformer.update(self._sim_dt) # Calculate the pose of the task frame in the root frame self._task_frame_pose_b[:, :3], self._task_frame_pose_b[:, 3:] = math_utils.subtract_frame_transforms( - wp.to_torch(self._asset.data.root_pos_w), - wp.to_torch(self._asset.data.root_quat_w), - wp.to_torch(self._task_frame_transformer.data.target_pos_w)[:, 0, :], - wp.to_torch(self._task_frame_transformer.data.target_quat_w)[:, 0, :], + self._asset.data.root_pos_w.torch, + self._asset.data.root_quat_w.torch, + self._task_frame_transformer.data.target_pos_w.torch[:, 0, :], + self._task_frame_transformer.data.target_quat_w.torch[:, 0, :], ) def _preprocess_actions(self, actions: torch.Tensor): diff --git a/source/isaaclab/isaaclab/envs/mdp/commands/pose_2d_command.py b/source/isaaclab/isaaclab/envs/mdp/commands/pose_2d_command.py index 1c181b3e32d7..37d7baf93155 100644 --- a/source/isaaclab/isaaclab/envs/mdp/commands/pose_2d_command.py +++ b/source/isaaclab/isaaclab/envs/mdp/commands/pose_2d_command.py @@ -11,7 +11,6 @@ from typing import TYPE_CHECKING import torch -import warp as wp from isaaclab.assets import Articulation from isaaclab.managers import CommandTerm @@ -84,11 +83,9 @@ def command(self) -> torch.Tensor: def _update_metrics(self): # logs data self.metrics["error_pos_2d"] = torch.linalg.norm( - self.pos_command_w[:, :2] - wp.to_torch(self.robot.data.root_pos_w)[:, :2], dim=1 - ) - self.metrics["error_heading"] = torch.abs( - wrap_to_pi(self.heading_command_w - wp.to_torch(self.robot.data.heading_w)) + self.pos_command_w[:, :2] - self.robot.data.root_pos_w.torch[:, :2], dim=1 ) + self.metrics["error_heading"] = torch.abs(wrap_to_pi(self.heading_command_w - self.robot.data.heading_w.torch)) def _resample_command(self, env_ids: Sequence[int]): # obtain env origins for the environments @@ -97,19 +94,19 @@ def _resample_command(self, env_ids: Sequence[int]): r = torch.empty(len(env_ids), device=self.device) self.pos_command_w[env_ids, 0] += r.uniform_(*self.cfg.ranges.pos_x) self.pos_command_w[env_ids, 1] += r.uniform_(*self.cfg.ranges.pos_y) - self.pos_command_w[env_ids, 2] += wp.to_torch(self.robot.data.default_root_pose)[env_ids, 2] + self.pos_command_w[env_ids, 2] += self.robot.data.default_root_pose.torch[env_ids, 2] if self.cfg.simple_heading: # set heading command to point towards target - target_vec = self.pos_command_w[env_ids] - wp.to_torch(self.robot.data.root_pos_w)[env_ids] + target_vec = self.pos_command_w[env_ids] - self.robot.data.root_pos_w.torch[env_ids] target_direction = torch.atan2(target_vec[:, 1], target_vec[:, 0]) flipped_target_direction = wrap_to_pi(target_direction + torch.pi) # compute errors to find the closest direction to the current heading # this is done to avoid the discontinuity at the -pi/pi boundary - curr_to_target = wrap_to_pi(target_direction - wp.to_torch(self.robot.data.heading_w)[env_ids]).abs() + curr_to_target = wrap_to_pi(target_direction - self.robot.data.heading_w.torch[env_ids]).abs() curr_to_flipped_target = wrap_to_pi( - flipped_target_direction - wp.to_torch(self.robot.data.heading_w)[env_ids] + flipped_target_direction - self.robot.data.heading_w.torch[env_ids] ).abs() # set the heading command to the closest direction @@ -124,9 +121,9 @@ def _resample_command(self, env_ids: Sequence[int]): def _update_command(self): """Re-target the position command to the current root state.""" - target_vec = self.pos_command_w - wp.to_torch(self.robot.data.root_pos_w)[:, :3] - self.pos_command_b[:] = quat_apply_inverse(yaw_quat(wp.to_torch(self.robot.data.root_quat_w)), target_vec) - self.heading_command_b[:] = wrap_to_pi(self.heading_command_w - wp.to_torch(self.robot.data.heading_w)) + target_vec = self.pos_command_w - self.robot.data.root_pos_w.torch[:, :3] + self.pos_command_b[:] = quat_apply_inverse(yaw_quat(self.robot.data.root_quat_w.torch), target_vec) + self.heading_command_b[:] = wrap_to_pi(self.heading_command_w - self.robot.data.heading_w.torch) def _set_debug_vis_impl(self, debug_vis: bool): # create markers if necessary for the first time @@ -186,19 +183,19 @@ def _resample_command(self, env_ids: Sequence[int]): self.terrain.terrain_levels[env_ids], self.terrain.terrain_types[env_ids], ids ] # offset the position command by the current root height - self.pos_command_w[env_ids, 2] += wp.to_torch(self.robot.data.default_root_pose)[env_ids, 2] + self.pos_command_w[env_ids, 2] += self.robot.data.default_root_pose.torch[env_ids, 2] if self.cfg.simple_heading: # set heading command to point towards target - target_vec = self.pos_command_w[env_ids] - wp.to_torch(self.robot.data.root_pos_w)[env_ids] + target_vec = self.pos_command_w[env_ids] - self.robot.data.root_pos_w.torch[env_ids] target_direction = torch.atan2(target_vec[:, 1], target_vec[:, 0]) flipped_target_direction = wrap_to_pi(target_direction + torch.pi) # compute errors to find the closest direction to the current heading # this is done to avoid the discontinuity at the -pi/pi boundary - curr_to_target = wrap_to_pi(target_direction - wp.to_torch(self.robot.data.heading_w)[env_ids]).abs() + curr_to_target = wrap_to_pi(target_direction - self.robot.data.heading_w.torch[env_ids]).abs() curr_to_flipped_target = wrap_to_pi( - flipped_target_direction - wp.to_torch(self.robot.data.heading_w)[env_ids] + flipped_target_direction - self.robot.data.heading_w.torch[env_ids] ).abs() # set the heading command to the closest direction diff --git a/source/isaaclab/isaaclab/envs/mdp/commands/pose_command.py b/source/isaaclab/isaaclab/envs/mdp/commands/pose_command.py index 33c58fd45f95..07f99a175155 100644 --- a/source/isaaclab/isaaclab/envs/mdp/commands/pose_command.py +++ b/source/isaaclab/isaaclab/envs/mdp/commands/pose_command.py @@ -11,7 +11,6 @@ from typing import TYPE_CHECKING import torch -import warp as wp from isaaclab.assets import Articulation from isaaclab.managers import CommandTerm @@ -94,8 +93,8 @@ def command(self) -> torch.Tensor: def _update_metrics(self): # transform command from base frame to simulation world frame self.pose_command_w[:, :3], self.pose_command_w[:, 3:] = combine_frame_transforms( - wp.to_torch(self.robot.data.root_pos_w), - wp.to_torch(self.robot.data.root_quat_w), + self.robot.data.root_pos_w.torch, + self.robot.data.root_quat_w.torch, self.pose_command_b[:, :3], self.pose_command_b[:, 3:], ) @@ -103,8 +102,8 @@ def _update_metrics(self): pos_error, rot_error = compute_pose_error( self.pose_command_w[:, :3], self.pose_command_w[:, 3:], - wp.to_torch(self.robot.data.body_pos_w)[:, self.body_idx], - wp.to_torch(self.robot.data.body_quat_w)[:, self.body_idx], + self.robot.data.body_pos_w.torch[:, self.body_idx], + self.robot.data.body_quat_w.torch[:, self.body_idx], ) self.metrics["position_error"] = torch.linalg.norm(pos_error, dim=-1) self.metrics["orientation_error"] = torch.linalg.norm(rot_error, dim=-1) @@ -153,5 +152,5 @@ def _debug_vis_callback(self, event): # -- goal pose self.goal_pose_visualizer.visualize(self.pose_command_w[:, :3], self.pose_command_w[:, 3:]) # -- current body pose - body_link_pose_w = wp.to_torch(self.robot.data.body_link_pose_w)[:, self.body_idx] + body_link_pose_w = self.robot.data.body_link_pose_w.torch[:, self.body_idx] self.current_pose_visualizer.visualize(body_link_pose_w[:, :3], body_link_pose_w[:, 3:7]) diff --git a/source/isaaclab/isaaclab/envs/mdp/commands/velocity_command.py b/source/isaaclab/isaaclab/envs/mdp/commands/velocity_command.py index eadc89af3af9..307ffc1c2868 100644 --- a/source/isaaclab/isaaclab/envs/mdp/commands/velocity_command.py +++ b/source/isaaclab/isaaclab/envs/mdp/commands/velocity_command.py @@ -12,7 +12,6 @@ from typing import TYPE_CHECKING import torch -import warp as wp import isaaclab.utils.math as math_utils from isaaclab.assets import Articulation @@ -118,11 +117,11 @@ def _update_metrics(self): max_command_step = max_command_time / self._env.step_dt # logs data self.metrics["error_vel_xy"] += ( - torch.linalg.norm(self.vel_command_b[:, :2] - wp.to_torch(self.robot.data.root_lin_vel_b)[:, :2], dim=-1) + torch.linalg.norm(self.vel_command_b[:, :2] - self.robot.data.root_lin_vel_b.torch[:, :2], dim=-1) / max_command_step ) self.metrics["error_vel_yaw"] += ( - torch.abs(self.vel_command_b[:, 2] - wp.to_torch(self.robot.data.root_ang_vel_b)[:, 2]) / max_command_step + torch.abs(self.vel_command_b[:, 2] - self.robot.data.root_ang_vel_b.torch[:, 2]) / max_command_step ) def _resample_command(self, env_ids: Sequence[int]): @@ -154,7 +153,7 @@ def _update_command(self): env_ids = self.is_heading_env.nonzero(as_tuple=False).flatten() # compute angular velocity heading_error = math_utils.wrap_to_pi( - self.heading_target[env_ids] - wp.to_torch(self.robot.data.heading_w)[env_ids] + self.heading_target[env_ids] - self.robot.data.heading_w.torch[env_ids] ) self.vel_command_b[env_ids, 2] = torch.clip( self.cfg.heading_control_stiffness * heading_error, @@ -191,12 +190,12 @@ def _debug_vis_callback(self, event): return # get marker location # -- base state - base_pos_w = wp.to_torch(self.robot.data.root_pos_w).clone() + base_pos_w = self.robot.data.root_pos_w.torch.clone() base_pos_w[:, 2] += 0.5 # -- resolve the scales and quaternions vel_des_arrow_scale, vel_des_arrow_quat = self._resolve_xy_velocity_to_arrow(self.command[:, :2]) vel_arrow_scale, vel_arrow_quat = self._resolve_xy_velocity_to_arrow( - wp.to_torch(self.robot.data.root_lin_vel_b)[:, :2] + self.robot.data.root_lin_vel_b.torch[:, :2] ) # display markers self.goal_vel_visualizer.visualize(base_pos_w, vel_des_arrow_quat, vel_des_arrow_scale) @@ -218,7 +217,7 @@ def _resolve_xy_velocity_to_arrow(self, xy_velocity: torch.Tensor) -> tuple[torc zeros = torch.zeros_like(heading_angle) arrow_quat = math_utils.quat_from_euler_xyz(zeros, zeros, heading_angle) # convert everything back from base to world frame - base_quat_w = wp.to_torch(self.robot.data.root_quat_w) + base_quat_w = self.robot.data.root_quat_w.torch arrow_quat = math_utils.quat_mul(base_quat_w, arrow_quat) return arrow_scale, arrow_quat diff --git a/source/isaaclab/isaaclab/envs/mdp/events.py b/source/isaaclab/isaaclab/envs/mdp/events.py index 9113eb6e4a7a..9483d2d2d0cd 100644 --- a/source/isaaclab/isaaclab/envs/mdp/events.py +++ b/source/isaaclab/isaaclab/envs/mdp/events.py @@ -483,9 +483,9 @@ def __call__( min_mass: float = 1e-6, ): if self.default_mass is None: - self.default_mass = wp.to_torch(self.asset.data.body_mass).clone() + self.default_mass = self.asset.data.body_mass.torch.clone() if self.default_inertia is None: - self.default_inertia = wp.to_torch(self.asset.data.body_inertia).clone() + self.default_inertia = self.asset.data.body_inertia.torch.clone() # resolve environment ids if env_ids is None: env_ids = torch.arange(env.scene.num_envs, device=self.asset.device, dtype=torch.int32) @@ -499,7 +499,7 @@ def __call__( body_ids = torch.tensor(self.asset_cfg.body_ids, dtype=torch.int32, device=self.asset.device) # get the current masses of the bodies (num_assets, num_bodies) - masses = wp.to_torch(self.asset.data.body_mass).clone() + masses = self.asset.data.body_mass.torch.clone() # apply randomization on default values # this is to make sure when calling the function multiple times, the randomization is applied on the @@ -524,7 +524,7 @@ def __call__( ratios = masses[env_ids[:, None], body_ids] / self.default_mass[env_ids[:, None], body_ids] # scale the inertia tensors by the the ratios # since mass randomization is done on default values, we can use the default inertia tensors - inertias = wp.to_torch(self.asset.data.body_inertia).clone() + inertias = self.asset.data.body_inertia.torch.clone() # inertia has shape: (num_envs, num_bodies, 9) for all assets inertias[env_ids[:, None], body_ids] = self.default_inertia[env_ids[:, None], body_ids] * ratios[..., None] # set the inertia tensors into the physics simulation @@ -626,7 +626,7 @@ def __call__( """ # store default inertia on first call for repeatable randomization if self.default_inertia is None: - self.default_inertia = wp.to_torch(self.asset.data.body_inertia).clone() + self.default_inertia = self.asset.data.body_inertia.torch.clone() # resolve environment ids if env_ids is None: @@ -714,7 +714,7 @@ def __call__( ): # store default CoM on first call for repeatable randomization if self.default_com is None: - self.default_com = wp.to_torch(self.asset.data.body_com_pose_b).clone() + self.default_com = self.asset.data.body_com_pose_b.torch.clone() # resolve environment ids if env_ids is None: @@ -1132,8 +1132,8 @@ def __init__(self, cfg: EventTermCfg, env: ManagerBasedEnv): self.asset_cfg: SceneEntityCfg = cfg.params["asset_cfg"] self.asset: RigidObject | Articulation = env.scene[self.asset_cfg.name] - self.default_joint_stiffness = wp.to_torch(self.asset.data.joint_stiffness).clone() - self.default_joint_damping = wp.to_torch(self.asset.data.joint_damping).clone() + self.default_joint_stiffness = self.asset.data.joint_stiffness.torch.clone() + self.default_joint_damping = self.asset.data.joint_damping.torch.clone() # For explicit actuators the sim-level stiffness/damping is zeroed out, so patch # the defaults with the actual actuator PD gains. @@ -1262,18 +1262,14 @@ def __init__(self, cfg: EventTermCfg, env: ManagerBasedEnv): self._backend = "newton" if "newton" in manager_name else "physx" # cache default values (common to both backends) - self.default_joint_friction_coeff = wp.to_torch(self.asset.data.joint_friction_coeff).clone() - self.default_joint_armature = wp.to_torch(self.asset.data.joint_armature).clone() - self.default_joint_pos_limits = wp.to_torch(self.asset.data.joint_pos_limits).clone() + self.default_joint_friction_coeff = self.asset.data.joint_friction_coeff.torch.clone() + self.default_joint_armature = self.asset.data.joint_armature.torch.clone() + self.default_joint_pos_limits = self.asset.data.joint_pos_limits.torch.clone() # cache dynamic/viscous friction (PhysX only - Newton only has static friction) if self._backend == "physx": - self.default_dynamic_joint_friction_coeff = wp.to_torch( - self.asset.data.joint_dynamic_friction_coeff - ).clone() - self.default_viscous_joint_friction_coeff = wp.to_torch( - self.asset.data.joint_viscous_friction_coeff - ).clone() + self.default_dynamic_joint_friction_coeff = (self.asset.data.joint_dynamic_friction_coeff.torch).clone() + self.default_viscous_joint_friction_coeff = (self.asset.data.joint_viscous_friction_coeff.torch).clone() # check for valid operation if cfg.params["operation"] == "scale": @@ -1381,7 +1377,7 @@ def __call__( # joint armature if armature_distribution_params is not None: armature = _randomize_prop_by_op( - wp.to_torch(self.asset.data.default_joint_armature).clone(), + self.asset.data.default_joint_armature.torch.clone(), armature_distribution_params, env_ids, joint_ids, @@ -1505,7 +1501,7 @@ def __call__( # stiffness if stiffness_distribution_params is not None: stiffness = _randomize_prop_by_op( - wp.to_torch(self.asset.data.fixed_tendon_stiffness).clone(), + self.asset.data.fixed_tendon_stiffness.torch.clone(), stiffness_distribution_params, env_ids, tendon_ids, @@ -1519,7 +1515,7 @@ def __call__( # damping if damping_distribution_params is not None: damping = _randomize_prop_by_op( - wp.to_torch(self.asset.data.fixed_tendon_damping).clone(), + self.asset.data.fixed_tendon_damping.torch.clone(), damping_distribution_params, env_ids, tendon_ids, @@ -1533,7 +1529,7 @@ def __call__( # limit stiffness if limit_stiffness_distribution_params is not None: limit_stiffness = _randomize_prop_by_op( - wp.to_torch(self.asset.data.fixed_tendon_limit_stiffness).clone(), + self.asset.data.fixed_tendon_limit_stiffness.torch.clone(), limit_stiffness_distribution_params, env_ids, tendon_ids, @@ -1546,7 +1542,7 @@ def __call__( # position limits if lower_limit_distribution_params is not None or upper_limit_distribution_params is not None: - limit = wp.to_torch(self.asset.data.fixed_tendon_pos_limits).clone() + limit = self.asset.data.fixed_tendon_pos_limits.torch.clone() # -- lower limit if lower_limit_distribution_params is not None: limit[..., 0] = _randomize_prop_by_op( @@ -1582,7 +1578,7 @@ def __call__( # rest length if rest_length_distribution_params is not None: rest_length = _randomize_prop_by_op( - wp.to_torch(self.asset.data.fixed_tendon_rest_length).clone(), + self.asset.data.fixed_tendon_rest_length.torch.clone(), rest_length_distribution_params, env_ids, tendon_ids, @@ -1596,7 +1592,7 @@ def __call__( # offset if offset_distribution_params is not None: offset = _randomize_prop_by_op( - wp.to_torch(self.asset.data.fixed_tendon_offset).clone(), + self.asset.data.fixed_tendon_offset.torch.clone(), offset_distribution_params, env_ids, tendon_ids, @@ -1668,7 +1664,7 @@ def push_by_setting_velocity( asset: RigidObject | Articulation = env.scene[asset_cfg.name] # velocities - vel_w = wp.to_torch(asset.data.root_vel_w)[env_ids] + vel_w = asset.data.root_vel_w.torch[env_ids] # sample random velocities range_list = [velocity_range.get(key, (0.0, 0.0)) for key in ["x", "y", "z", "roll", "pitch", "yaw"]] ranges = torch.tensor(range_list, device=asset.device) @@ -1700,8 +1696,8 @@ def reset_root_state_uniform( # extract the used quantities (to enable type-hinting) asset: RigidObject | Articulation = env.scene[asset_cfg.name] # get default root state - default_root_pose = wp.to_torch(asset.data.default_root_pose)[env_ids].clone() - default_root_vel = wp.to_torch(asset.data.default_root_vel)[env_ids].clone() + default_root_pose = asset.data.default_root_pose.torch[env_ids].clone() + default_root_vel = asset.data.default_root_vel.torch[env_ids].clone() # poses range_list = [pose_range.get(key, (0.0, 0.0)) for key in ["x", "y", "z", "roll", "pitch", "yaw"]] @@ -1753,8 +1749,8 @@ def reset_root_state_with_random_orientation( # extract the used quantities (to enable type-hinting) asset: RigidObject | Articulation = env.scene[asset_cfg.name] # get default root state - default_root_pose = wp.to_torch(asset.data.default_root_pose)[env_ids].clone() - default_root_vel = wp.to_torch(asset.data.default_root_vel)[env_ids].clone() + default_root_pose = asset.data.default_root_pose.torch[env_ids].clone() + default_root_vel = asset.data.default_root_vel.torch[env_ids].clone() # poses range_list = [pose_range.get(key, (0.0, 0.0)) for key in ["x", "y", "z"]] @@ -1821,7 +1817,7 @@ def reset_root_state_from_terrain( # sample random valid poses ids = torch.randint(0, valid_positions.shape[2], size=(len(env_ids),), device=env.device) positions = valid_positions[terrain.terrain_levels[env_ids], terrain.terrain_types[env_ids], ids] - positions += wp.to_torch(asset.data.default_root_pose)[env_ids, :3] + positions += asset.data.default_root_pose.torch[env_ids, :3] # sample random orientations range_list = [pose_range.get(key, (0.0, 0.0)) for key in ["roll", "pitch", "yaw"]] @@ -1836,7 +1832,7 @@ def reset_root_state_from_terrain( ranges = torch.tensor(range_list, device=asset.device) rand_samples = math_utils.sample_uniform(ranges[:, 0], ranges[:, 1], (len(env_ids), 6), device=asset.device) - velocities = wp.to_torch(asset.data.default_root_vel)[env_ids] + rand_samples + velocities = asset.data.default_root_vel.torch[env_ids] + rand_samples # set into the physics simulation asset.write_root_pose_to_sim_index(root_pose=torch.cat([positions, orientations], dim=-1), env_ids=env_ids) @@ -1865,18 +1861,18 @@ def reset_joints_by_scale( iter_env_ids = env_ids # get default joint state - joint_pos = wp.to_torch(asset.data.default_joint_pos)[iter_env_ids, asset_cfg.joint_ids].clone() - joint_vel = wp.to_torch(asset.data.default_joint_vel)[iter_env_ids, asset_cfg.joint_ids].clone() + joint_pos = asset.data.default_joint_pos.torch[iter_env_ids, asset_cfg.joint_ids].clone() + joint_vel = asset.data.default_joint_vel.torch[iter_env_ids, asset_cfg.joint_ids].clone() # scale these values randomly joint_pos *= math_utils.sample_uniform(*position_range, joint_pos.shape, joint_pos.device) joint_vel *= math_utils.sample_uniform(*velocity_range, joint_vel.shape, joint_vel.device) # clamp joint pos to limits - joint_pos_limits = wp.to_torch(asset.data.soft_joint_pos_limits)[iter_env_ids, asset_cfg.joint_ids] + joint_pos_limits = asset.data.soft_joint_pos_limits.torch[iter_env_ids, asset_cfg.joint_ids] joint_pos = joint_pos.clamp_(joint_pos_limits[..., 0], joint_pos_limits[..., 1]) # clamp joint vel to limits - joint_vel_limits = wp.to_torch(asset.data.soft_joint_vel_limits)[iter_env_ids, asset_cfg.joint_ids] + joint_vel_limits = asset.data.soft_joint_vel_limits.torch[iter_env_ids, asset_cfg.joint_ids] joint_vel = joint_vel.clamp_(-joint_vel_limits, joint_vel_limits) # set into the physics simulation @@ -1906,18 +1902,18 @@ def reset_joints_by_offset( iter_env_ids = env_ids # get default joint state - joint_pos = wp.to_torch(asset.data.default_joint_pos)[iter_env_ids, asset_cfg.joint_ids].clone() - joint_vel = wp.to_torch(asset.data.default_joint_vel)[iter_env_ids, asset_cfg.joint_ids].clone() + joint_pos = asset.data.default_joint_pos.torch[iter_env_ids, asset_cfg.joint_ids].clone() + joint_vel = asset.data.default_joint_vel.torch[iter_env_ids, asset_cfg.joint_ids].clone() # bias these values randomly joint_pos += math_utils.sample_uniform(*position_range, joint_pos.shape, joint_pos.device) joint_vel += math_utils.sample_uniform(*velocity_range, joint_vel.shape, joint_vel.device) # clamp joint pos to limits - joint_pos_limits = wp.to_torch(asset.data.soft_joint_pos_limits)[iter_env_ids, asset_cfg.joint_ids] + joint_pos_limits = asset.data.soft_joint_pos_limits.torch[iter_env_ids, asset_cfg.joint_ids] joint_pos = joint_pos.clamp_(joint_pos_limits[..., 0], joint_pos_limits[..., 1]) # clamp joint vel to limits - joint_vel_limits = wp.to_torch(asset.data.soft_joint_vel_limits)[iter_env_ids, asset_cfg.joint_ids] + joint_vel_limits = asset.data.soft_joint_vel_limits.torch[iter_env_ids, asset_cfg.joint_ids] joint_vel = joint_vel.clamp_(-joint_vel_limits, joint_vel_limits) # set into the physics simulation @@ -1947,7 +1943,7 @@ def reset_nodal_state_uniform( # extract the used quantities (to enable type-hinting) asset: DeformableObject = env.scene[asset_cfg.name] # get default root state - nodal_state = wp.to_torch(asset.data.default_nodal_state_w)[env_ids].clone() + nodal_state = asset.data.default_nodal_state_w.torch[env_ids].clone() # position range_list = [position_range.get(key, (0.0, 0.0)) for key in ["x", "y", "z"]] @@ -1978,8 +1974,8 @@ def reset_scene_to_default(env: ManagerBasedEnv, env_ids: torch.Tensor, reset_jo # rigid bodies for rigid_object in env.scene.rigid_objects.values(): # obtain default and deal with the offset for env origins - default_root_pose = wp.to_torch(rigid_object.data.default_root_pose)[env_ids].clone() - default_root_vel = wp.to_torch(rigid_object.data.default_root_vel)[env_ids].clone() + default_root_pose = rigid_object.data.default_root_pose.torch[env_ids].clone() + default_root_vel = rigid_object.data.default_root_vel.torch[env_ids].clone() default_root_pose[:, :3] += env.scene.env_origins[env_ids] # set into the physics simulation rigid_object.write_root_pose_to_sim_index(root_pose=default_root_pose, env_ids=env_ids) @@ -1987,15 +1983,15 @@ def reset_scene_to_default(env: ManagerBasedEnv, env_ids: torch.Tensor, reset_jo # articulations for articulation_asset in env.scene.articulations.values(): # obtain default and deal with the offset for env origins - default_root_pose = wp.to_torch(articulation_asset.data.default_root_pose)[env_ids].clone() - default_root_vel = wp.to_torch(articulation_asset.data.default_root_vel)[env_ids].clone() + default_root_pose = articulation_asset.data.default_root_pose.torch[env_ids].clone() + default_root_vel = articulation_asset.data.default_root_vel.torch[env_ids].clone() default_root_pose[:, :3] += env.scene.env_origins[env_ids] # set into the physics simulation articulation_asset.write_root_pose_to_sim_index(root_pose=default_root_pose, env_ids=env_ids) articulation_asset.write_root_velocity_to_sim_index(root_velocity=default_root_vel, env_ids=env_ids) # obtain default joint positions - default_joint_pos = wp.to_torch(articulation_asset.data.default_joint_pos)[env_ids].clone() - default_joint_vel = wp.to_torch(articulation_asset.data.default_joint_vel)[env_ids].clone() + default_joint_pos = articulation_asset.data.default_joint_pos.torch[env_ids].clone() + default_joint_vel = articulation_asset.data.default_joint_vel.torch[env_ids].clone() # set into the physics simulation articulation_asset.write_joint_position_to_sim_index(position=default_joint_pos, env_ids=env_ids) articulation_asset.write_joint_velocity_to_sim_index(velocity=default_joint_vel, env_ids=env_ids) @@ -2006,7 +2002,7 @@ def reset_scene_to_default(env: ManagerBasedEnv, env_ids: torch.Tensor, reset_jo # deformable objects for deformable_object in env.scene.deformable_objects.values(): # obtain default and set into the physics simulation - nodal_state = wp.to_torch(deformable_object.data.default_nodal_state_w)[env_ids].clone() + nodal_state = deformable_object.data.default_nodal_state_w.torch[env_ids].clone() deformable_object.write_nodal_state_to_sim(nodal_state, env_ids=env_ids) diff --git a/source/isaaclab/isaaclab/envs/mdp/observations.py b/source/isaaclab/isaaclab/envs/mdp/observations.py index d65bd264f7ed..56d369d18718 100644 --- a/source/isaaclab/isaaclab/envs/mdp/observations.py +++ b/source/isaaclab/isaaclab/envs/mdp/observations.py @@ -14,7 +14,6 @@ from typing import TYPE_CHECKING import torch -import warp as wp import isaaclab.utils.math as math_utils from isaaclab.managers import SceneEntityCfg @@ -46,7 +45,7 @@ def base_pos_z(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg( """Root height in the simulation world frame.""" # extract the used quantities (to enable type-hinting) asset: Articulation = env.scene[asset_cfg.name] - return wp.to_torch(asset.data.root_pos_w)[:, 2].unsqueeze(-1) + return asset.data.root_pos_w.torch[:, 2].unsqueeze(-1) @generic_io_descriptor( @@ -56,7 +55,7 @@ def base_lin_vel(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCf """Root linear velocity in the asset's root frame.""" # extract the used quantities (to enable type-hinting) asset: RigidObject = env.scene[asset_cfg.name] - return wp.to_torch(asset.data.root_lin_vel_b) + return asset.data.root_lin_vel_b.torch @generic_io_descriptor( @@ -66,7 +65,7 @@ def base_ang_vel(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCf """Root angular velocity in the asset's root frame.""" # extract the used quantities (to enable type-hinting) asset: RigidObject = env.scene[asset_cfg.name] - return wp.to_torch(asset.data.root_ang_vel_b) + return asset.data.root_ang_vel_b.torch @generic_io_descriptor( @@ -76,7 +75,7 @@ def projected_gravity(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEnt """Gravity projection on the asset's root frame.""" # extract the used quantities (to enable type-hinting) asset: RigidObject = env.scene[asset_cfg.name] - return wp.to_torch(asset.data.projected_gravity_b) + return asset.data.projected_gravity_b.torch @generic_io_descriptor( @@ -86,7 +85,7 @@ def root_pos_w(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg( """Asset root position in the environment frame.""" # extract the used quantities (to enable type-hinting) asset: RigidObject = env.scene[asset_cfg.name] - return wp.to_torch(asset.data.root_pos_w) - env.scene.env_origins + return asset.data.root_pos_w.torch - env.scene.env_origins @generic_io_descriptor( @@ -104,7 +103,7 @@ def root_quat_w( # extract the used quantities (to enable type-hinting) asset: RigidObject = env.scene[asset_cfg.name] - quat = wp.to_torch(asset.data.root_quat_w) + quat = asset.data.root_quat_w.torch # make the quaternion real-part positive if configured return math_utils.quat_unique(quat) if make_quat_unique else quat @@ -116,7 +115,7 @@ def root_lin_vel_w(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntity """Asset root linear velocity in the environment frame.""" # extract the used quantities (to enable type-hinting) asset: RigidObject = env.scene[asset_cfg.name] - return wp.to_torch(asset.data.root_lin_vel_w) + return asset.data.root_lin_vel_w.torch @generic_io_descriptor( @@ -126,7 +125,7 @@ def root_ang_vel_w(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntity """Asset root angular velocity in the environment frame.""" # extract the used quantities (to enable type-hinting) asset: RigidObject = env.scene[asset_cfg.name] - return wp.to_torch(asset.data.root_ang_vel_w) + return asset.data.root_ang_vel_w.torch """ @@ -155,7 +154,7 @@ def body_pose_w( asset: Articulation = env.scene[asset_cfg.name] # access the body poses in world frame - pose = wp.to_torch(asset.data.body_pose_w)[:, asset_cfg.body_ids, :7] + pose = asset.data.body_pose_w.torch[:, asset_cfg.body_ids, :7] if isinstance(asset_cfg.body_ids, (slice, int)): pose = pose.clone() # if slice or int, make a copy to avoid modifying original data pose[..., :3] = pose[..., :3] - env.scene.env_origins.unsqueeze(1) @@ -182,8 +181,8 @@ def body_projected_gravity_b( # extract the used quantities (to enable type-hinting) asset: Articulation = env.scene[asset_cfg.name] - body_quat = wp.to_torch(asset.data.body_quat_w)[:, asset_cfg.body_ids] - gravity_dir = wp.to_torch(asset.data.GRAVITY_VEC_W).unsqueeze(1) + body_quat = asset.data.body_quat_w.torch[:, asset_cfg.body_ids] + gravity_dir = asset.data.GRAVITY_VEC_W.torch.unsqueeze(1) return math_utils.quat_apply_inverse(body_quat, gravity_dir).view(env.num_envs, -1) @@ -202,7 +201,7 @@ def joint_pos(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg(" """ # extract the used quantities (to enable type-hinting) asset: Articulation = env.scene[asset_cfg.name] - return wp.to_torch(asset.data.joint_pos)[:, asset_cfg.joint_ids] + return asset.data.joint_pos.torch[:, asset_cfg.joint_ids] @generic_io_descriptor( @@ -218,8 +217,7 @@ def joint_pos_rel(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityC # extract the used quantities (to enable type-hinting) asset: Articulation = env.scene[asset_cfg.name] return ( - wp.to_torch(asset.data.joint_pos)[:, asset_cfg.joint_ids] - - wp.to_torch(asset.data.default_joint_pos)[:, asset_cfg.joint_ids] + asset.data.joint_pos.torch[:, asset_cfg.joint_ids] - asset.data.default_joint_pos.torch[:, asset_cfg.joint_ids] ) @@ -234,9 +232,9 @@ def joint_pos_limit_normalized( # extract the used quantities (to enable type-hinting) asset: Articulation = env.scene[asset_cfg.name] return math_utils.scale_transform( - wp.to_torch(asset.data.joint_pos)[:, asset_cfg.joint_ids], - wp.to_torch(asset.data.soft_joint_pos_limits)[:, asset_cfg.joint_ids, 0], - wp.to_torch(asset.data.soft_joint_pos_limits)[:, asset_cfg.joint_ids, 1], + asset.data.joint_pos.torch[:, asset_cfg.joint_ids], + asset.data.soft_joint_pos_limits.torch[:, asset_cfg.joint_ids, 0], + asset.data.soft_joint_pos_limits.torch[:, asset_cfg.joint_ids, 1], ) @@ -250,7 +248,7 @@ def joint_vel(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg(" """ # extract the used quantities (to enable type-hinting) asset: Articulation = env.scene[asset_cfg.name] - return wp.to_torch(asset.data.joint_vel)[:, asset_cfg.joint_ids] + return asset.data.joint_vel.torch[:, asset_cfg.joint_ids] @generic_io_descriptor( @@ -266,8 +264,7 @@ def joint_vel_rel(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityC # extract the used quantities (to enable type-hinting) asset: Articulation = env.scene[asset_cfg.name] return ( - wp.to_torch(asset.data.joint_vel)[:, asset_cfg.joint_ids] - - wp.to_torch(asset.data.default_joint_vel)[:, asset_cfg.joint_ids] + asset.data.joint_vel.torch[:, asset_cfg.joint_ids] - asset.data.default_joint_vel.torch[:, asset_cfg.joint_ids] ) @@ -288,7 +285,7 @@ def joint_effort(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCf """ # extract the used quantities (to enable type-hinting) asset: Articulation = env.scene[asset_cfg.name] - return wp.to_torch(asset.data.applied_torque)[:, asset_cfg.joint_ids] + return asset.data.applied_torque.torch[:, asset_cfg.joint_ids] """ @@ -304,7 +301,7 @@ def height_scan(env: ManagerBasedEnv, sensor_cfg: SceneEntityCfg, offset: float # extract the used quantities (to enable type-hinting) sensor: RayCaster = env.scene.sensors[sensor_cfg.name] # height scan: height = sensor_height - hit_point_z - offset - return wp.to_torch(sensor.data.pos_w)[:, 2].unsqueeze(1) - wp.to_torch(sensor.data.ray_hits_w)[..., 2] - offset + return sensor.data.pos_w.torch[:, 2].unsqueeze(1) - sensor.data.ray_hits_w.torch[..., 2] - offset def body_incoming_wrench(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg) -> torch.Tensor: @@ -315,7 +312,7 @@ def body_incoming_wrench(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg) -> tor # extract the used quantities (to enable type-hinting) asset: Articulation = env.scene[asset_cfg.name] # obtain the link incoming forces in world frame - body_incoming_joint_wrench_b = wp.to_torch(asset.data.body_incoming_joint_wrench_b)[:, asset_cfg.body_ids] + body_incoming_joint_wrench_b = asset.data.body_incoming_joint_wrench_b.torch[:, asset_cfg.body_ids] return body_incoming_joint_wrench_b.view(env.num_envs, -1) @@ -330,7 +327,7 @@ def pva_orientation(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntit Orientation in the world frame in (x, y, z, w) quaternion form. Shape is (num_envs, 4). """ asset: Pva = env.scene[asset_cfg.name] - return wp.to_torch(asset.data.quat_w) + return asset.data.quat_w.torch def pva_projected_gravity(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("pva")) -> torch.Tensor: @@ -344,7 +341,7 @@ def pva_projected_gravity(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = Scen Gravity projected on sensor frame, shape of torch.tensor is (num_env, 3). """ asset: Pva = env.scene[asset_cfg.name] - return wp.to_torch(asset.data.projected_gravity_b) + return asset.data.projected_gravity_b.torch def imu_ang_vel(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("imu")) -> torch.Tensor: @@ -358,7 +355,7 @@ def imu_ang_vel(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg The angular velocity [rad/s] in the sensor frame. Shape is (num_envs, 3). """ asset: Imu = env.scene[asset_cfg.name] - return wp.to_torch(asset.data.ang_vel_b) + return asset.data.ang_vel_b.torch def imu_lin_acc(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("imu")) -> torch.Tensor: @@ -372,7 +369,7 @@ def imu_lin_acc(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg The linear acceleration [m/s^2] in the sensor frame. Shape is (num_envs, 3). """ asset: Imu = env.scene[asset_cfg.name] - return wp.to_torch(asset.data.lin_acc_b) + return asset.data.lin_acc_b.torch def image( diff --git a/source/isaaclab/isaaclab/envs/mdp/rewards.py b/source/isaaclab/isaaclab/envs/mdp/rewards.py index 74bea7ee7861..4645e4fd805e 100644 --- a/source/isaaclab/isaaclab/envs/mdp/rewards.py +++ b/source/isaaclab/isaaclab/envs/mdp/rewards.py @@ -14,7 +14,6 @@ from typing import TYPE_CHECKING import torch -import warp as wp from isaaclab.managers import SceneEntityCfg from isaaclab.managers.manager_base import ManagerTermBase @@ -79,14 +78,14 @@ def lin_vel_z_l2(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg = SceneEntity """Penalize z-axis base linear velocity using L2 squared kernel.""" # extract the used quantities (to enable type-hinting) asset: RigidObject = env.scene[asset_cfg.name] - return torch.square(wp.to_torch(asset.data.root_lin_vel_b)[:, 2]) + return torch.square(asset.data.root_lin_vel_b.torch[:, 2]) def ang_vel_xy_l2(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor: """Penalize xy-axis base angular velocity using L2 squared kernel.""" # extract the used quantities (to enable type-hinting) asset: RigidObject = env.scene[asset_cfg.name] - return torch.sum(torch.square(wp.to_torch(asset.data.root_ang_vel_b)[:, :2]), dim=1) + return torch.sum(torch.square(asset.data.root_ang_vel_b.torch[:, :2]), dim=1) def flat_orientation_l2(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor: @@ -96,7 +95,7 @@ def flat_orientation_l2(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg = Scen """ # extract the used quantities (to enable type-hinting) asset: RigidObject = env.scene[asset_cfg.name] - return torch.sum(torch.square(wp.to_torch(asset.data.projected_gravity_b)[:, :2]), dim=1) + return torch.sum(torch.square(asset.data.projected_gravity_b.torch[:, :2]), dim=1) def base_height_l2( @@ -116,18 +115,18 @@ def base_height_l2( if sensor_cfg is not None: sensor: RayCaster = env.scene[sensor_cfg.name] # Adjust the target height using the sensor data - adjusted_target_height = target_height + torch.mean(wp.to_torch(sensor.data.ray_hits_w)[..., 2], dim=1) + adjusted_target_height = target_height + torch.mean(sensor.data.ray_hits_w.torch[..., 2], dim=1) else: # Use the provided target height directly for flat terrain adjusted_target_height = target_height # Compute the L2 squared penalty - return torch.square(wp.to_torch(asset.data.root_pos_w)[:, 2] - adjusted_target_height) + return torch.square(asset.data.root_pos_w.torch[:, 2] - adjusted_target_height) def body_lin_acc_l2(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor: """Penalize the linear acceleration of bodies using L2-kernel.""" asset: Articulation = env.scene[asset_cfg.name] - return torch.sum(torch.linalg.norm(wp.to_torch(asset.data.body_lin_acc_w)[:, asset_cfg.body_ids, :], dim=-1), dim=1) + return torch.sum(torch.linalg.norm(asset.data.body_lin_acc_w.torch[:, asset_cfg.body_ids, :], dim=-1), dim=1) """ @@ -144,14 +143,14 @@ def joint_torques_l2(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg = SceneEn """ # extract the used quantities (to enable type-hinting) asset: Articulation = env.scene[asset_cfg.name] - return torch.sum(torch.square(wp.to_torch(asset.data.applied_torque)[:, asset_cfg.joint_ids]), dim=1) + return torch.sum(torch.square(asset.data.applied_torque.torch[:, asset_cfg.joint_ids]), dim=1) def joint_vel_l1(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg) -> torch.Tensor: """Penalize joint velocities on the articulation using an L1-kernel.""" # extract the used quantities (to enable type-hinting) asset: Articulation = env.scene[asset_cfg.name] - return torch.sum(torch.abs(wp.to_torch(asset.data.joint_vel)[:, asset_cfg.joint_ids]), dim=1) + return torch.sum(torch.abs(asset.data.joint_vel.torch[:, asset_cfg.joint_ids]), dim=1) def joint_vel_l2(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor: @@ -163,7 +162,7 @@ def joint_vel_l2(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg = SceneEntity """ # extract the used quantities (to enable type-hinting) asset: Articulation = env.scene[asset_cfg.name] - return torch.sum(torch.square(wp.to_torch(asset.data.joint_vel)[:, asset_cfg.joint_ids]), dim=1) + return torch.sum(torch.square(asset.data.joint_vel.torch[:, asset_cfg.joint_ids]), dim=1) def joint_acc_l2(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor: @@ -175,7 +174,7 @@ def joint_acc_l2(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg = SceneEntity """ # extract the used quantities (to enable type-hinting) asset: Articulation = env.scene[asset_cfg.name] - return torch.sum(torch.square(wp.to_torch(asset.data.joint_acc)[:, asset_cfg.joint_ids]), dim=1) + return torch.sum(torch.square(asset.data.joint_acc.torch[:, asset_cfg.joint_ids]), dim=1) def joint_deviation_l1(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor: @@ -184,8 +183,7 @@ def joint_deviation_l1(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg = Scene asset: Articulation = env.scene[asset_cfg.name] # compute out of limits constraints angle = ( - wp.to_torch(asset.data.joint_pos)[:, asset_cfg.joint_ids] - - wp.to_torch(asset.data.default_joint_pos)[:, asset_cfg.joint_ids] + asset.data.joint_pos.torch[:, asset_cfg.joint_ids] - asset.data.default_joint_pos.torch[:, asset_cfg.joint_ids] ) return torch.sum(torch.abs(angle), dim=1) @@ -199,12 +197,12 @@ def joint_pos_limits(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg = SceneEn asset: Articulation = env.scene[asset_cfg.name] # compute out of limits constraints out_of_limits = -( - wp.to_torch(asset.data.joint_pos)[:, asset_cfg.joint_ids] - - wp.to_torch(asset.data.soft_joint_pos_limits)[:, asset_cfg.joint_ids, 0] + asset.data.joint_pos.torch[:, asset_cfg.joint_ids] + - asset.data.soft_joint_pos_limits.torch[:, asset_cfg.joint_ids, 0] ).clip(max=0.0) out_of_limits += ( - wp.to_torch(asset.data.joint_pos)[:, asset_cfg.joint_ids] - - wp.to_torch(asset.data.soft_joint_pos_limits)[:, asset_cfg.joint_ids, 1] + asset.data.joint_pos.torch[:, asset_cfg.joint_ids] + - asset.data.soft_joint_pos_limits.torch[:, asset_cfg.joint_ids, 1] ).clip(min=0.0) return torch.sum(out_of_limits, dim=1) @@ -223,8 +221,8 @@ def joint_vel_limits( asset: Articulation = env.scene[asset_cfg.name] # compute out of limits constraints out_of_limits = ( - torch.abs(wp.to_torch(asset.data.joint_vel)[:, asset_cfg.joint_ids]) - - wp.to_torch(asset.data.soft_joint_vel_limits)[:, asset_cfg.joint_ids] * soft_ratio + torch.abs(asset.data.joint_vel.torch[:, asset_cfg.joint_ids]) + - asset.data.soft_joint_vel_limits.torch[:, asset_cfg.joint_ids] * soft_ratio ) # clip to max error = 1 rad/s per joint to avoid huge penalties out_of_limits = out_of_limits.clip_(min=0.0, max=1.0) @@ -250,8 +248,8 @@ def applied_torque_limits(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg = Sc # compute out of limits constraints # TODO: We need to fix this to support implicit joints. out_of_limits = torch.abs( - wp.to_torch(asset.data.applied_torque)[:, asset_cfg.joint_ids] - - wp.to_torch(asset.data.computed_torque)[:, asset_cfg.joint_ids] + asset.data.applied_torque.torch[:, asset_cfg.joint_ids] + - asset.data.computed_torque.torch[:, asset_cfg.joint_ids] ) return torch.sum(out_of_limits, dim=1) @@ -276,7 +274,7 @@ def undesired_contacts(env: ManagerBasedRLEnv, threshold: float, sensor_cfg: Sce # extract the used quantities (to enable type-hinting) contact_sensor: ContactSensor = env.scene.sensors[sensor_cfg.name] # check if contact force is above threshold - net_contact_forces = wp.to_torch(contact_sensor.data.net_forces_w_history) + net_contact_forces = contact_sensor.data.net_forces_w_history.torch is_contact = ( torch.max(torch.linalg.norm(net_contact_forces[:, :, sensor_cfg.body_ids], dim=-1), dim=1)[0] > threshold ) @@ -288,7 +286,7 @@ def desired_contacts(env, sensor_cfg: SceneEntityCfg, threshold: float = 1.0) -> """Penalize if none of the desired contacts are present.""" contact_sensor: ContactSensor = env.scene.sensors[sensor_cfg.name] contacts = ( - wp.to_torch(contact_sensor.data.net_forces_w_history)[:, :, sensor_cfg.body_ids, :].norm(dim=-1).max(dim=1)[0] + contact_sensor.data.net_forces_w_history.torch[:, :, sensor_cfg.body_ids, :].norm(dim=-1).max(dim=1)[0] > threshold ) zero_contact = (~contacts).all(dim=1) @@ -299,7 +297,7 @@ def contact_forces(env: ManagerBasedRLEnv, threshold: float, sensor_cfg: SceneEn """Penalize contact forces as the amount of violations of the net contact force.""" # extract the used quantities (to enable type-hinting) contact_sensor: ContactSensor = env.scene.sensors[sensor_cfg.name] - net_contact_forces = wp.to_torch(contact_sensor.data.net_forces_w_history) + net_contact_forces = contact_sensor.data.net_forces_w_history.torch # compute the violation violation = ( torch.max(torch.linalg.norm(net_contact_forces[:, :, sensor_cfg.body_ids], dim=-1), dim=1)[0] - threshold @@ -321,9 +319,7 @@ def track_lin_vel_xy_exp( asset: RigidObject = env.scene[asset_cfg.name] # compute the error lin_vel_error = torch.sum( - torch.square( - env.command_manager.get_command(command_name)[:, :2] - wp.to_torch(asset.data.root_lin_vel_b)[:, :2] - ), + torch.square(env.command_manager.get_command(command_name)[:, :2] - asset.data.root_lin_vel_b.torch[:, :2]), dim=1, ) return torch.exp(-lin_vel_error / std**2) @@ -337,6 +333,6 @@ def track_ang_vel_z_exp( asset: RigidObject = env.scene[asset_cfg.name] # compute the error ang_vel_error = torch.square( - env.command_manager.get_command(command_name)[:, 2] - wp.to_torch(asset.data.root_ang_vel_b)[:, 2] + env.command_manager.get_command(command_name)[:, 2] - asset.data.root_ang_vel_b.torch[:, 2] ) return torch.exp(-ang_vel_error / std**2) diff --git a/source/isaaclab/isaaclab/envs/mdp/terminations.py b/source/isaaclab/isaaclab/envs/mdp/terminations.py index 1936f5dd50b6..d717fb0dfcd0 100644 --- a/source/isaaclab/isaaclab/envs/mdp/terminations.py +++ b/source/isaaclab/isaaclab/envs/mdp/terminations.py @@ -14,7 +14,6 @@ from typing import TYPE_CHECKING import torch -import warp as wp from isaaclab.managers import SceneEntityCfg @@ -58,7 +57,7 @@ def bad_orientation( """ # extract the used quantities (to enable type-hinting) asset: RigidObject = env.scene[asset_cfg.name] - return torch.acos(-wp.to_torch(asset.data.projected_gravity_b)[:, 2]).abs() > limit_angle + return torch.acos(-asset.data.projected_gravity_b.torch[:, 2]).abs() > limit_angle def root_height_below_minimum( @@ -71,7 +70,7 @@ def root_height_below_minimum( """ # extract the used quantities (to enable type-hinting) asset: RigidObject = env.scene[asset_cfg.name] - return wp.to_torch(asset.data.root_pos_w)[:, 2] < minimum_height + return asset.data.root_pos_w.torch[:, 2] < minimum_height """ @@ -86,9 +85,9 @@ def joint_pos_out_of_limit(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg = S if asset_cfg.joint_ids is None: asset_cfg.joint_ids = slice(None) - limits = wp.to_torch(asset.data.soft_joint_pos_limits)[:, asset_cfg.joint_ids] - out_of_upper_limits = torch.any(wp.to_torch(asset.data.joint_pos)[:, asset_cfg.joint_ids] > limits[..., 1], dim=1) - out_of_lower_limits = torch.any(wp.to_torch(asset.data.joint_pos)[:, asset_cfg.joint_ids] < limits[..., 0], dim=1) + limits = asset.data.soft_joint_pos_limits.torch[:, asset_cfg.joint_ids] + out_of_upper_limits = torch.any(asset.data.joint_pos.torch[:, asset_cfg.joint_ids] > limits[..., 1], dim=1) + out_of_lower_limits = torch.any(asset.data.joint_pos.torch[:, asset_cfg.joint_ids] < limits[..., 0], dim=1) return torch.logical_or(out_of_upper_limits, out_of_lower_limits) @@ -105,8 +104,8 @@ def joint_pos_out_of_manual_limit( if asset_cfg.joint_ids is None: asset_cfg.joint_ids = slice(None) # compute any violations - out_of_upper_limits = torch.any(wp.to_torch(asset.data.joint_pos)[:, asset_cfg.joint_ids] > bounds[1], dim=1) - out_of_lower_limits = torch.any(wp.to_torch(asset.data.joint_pos)[:, asset_cfg.joint_ids] < bounds[0], dim=1) + out_of_upper_limits = torch.any(asset.data.joint_pos.torch[:, asset_cfg.joint_ids] > bounds[1], dim=1) + out_of_lower_limits = torch.any(asset.data.joint_pos.torch[:, asset_cfg.joint_ids] < bounds[0], dim=1) return torch.logical_or(out_of_upper_limits, out_of_lower_limits) @@ -115,9 +114,9 @@ def joint_vel_out_of_limit(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg = S # extract the used quantities (to enable type-hinting) asset: Articulation = env.scene[asset_cfg.name] # compute any violations - limits = wp.to_torch(asset.data.soft_joint_vel_limits) + limits = asset.data.soft_joint_vel_limits.torch return torch.any( - torch.abs(wp.to_torch(asset.data.joint_vel)[:, asset_cfg.joint_ids]) > limits[:, asset_cfg.joint_ids], dim=1 + torch.abs(asset.data.joint_vel.torch[:, asset_cfg.joint_ids]) > limits[:, asset_cfg.joint_ids], dim=1 ) @@ -128,7 +127,7 @@ def joint_vel_out_of_manual_limit( # extract the used quantities (to enable type-hinting) asset: Articulation = env.scene[asset_cfg.name] # compute any violations - return torch.any(torch.abs(wp.to_torch(asset.data.joint_vel)[:, asset_cfg.joint_ids]) > max_velocity, dim=1) + return torch.any(torch.abs(asset.data.joint_vel.torch[:, asset_cfg.joint_ids]) > max_velocity, dim=1) def joint_effort_out_of_limit( @@ -144,8 +143,8 @@ def joint_effort_out_of_limit( asset: Articulation = env.scene[asset_cfg.name] # check if any joint effort is out of limit out_of_limits = ~torch.isclose( - wp.to_torch(asset.data.computed_torque)[:, asset_cfg.joint_ids], - wp.to_torch(asset.data.applied_torque)[:, asset_cfg.joint_ids], + asset.data.computed_torque.torch[:, asset_cfg.joint_ids], + asset.data.applied_torque.torch[:, asset_cfg.joint_ids], ) return torch.any(out_of_limits, dim=1) @@ -159,7 +158,7 @@ def illegal_contact(env: ManagerBasedRLEnv, threshold: float, sensor_cfg: SceneE """Terminate when the contact force on the sensor exceeds the force threshold.""" # extract the used quantities (to enable type-hinting) contact_sensor: ContactSensor = env.scene.sensors[sensor_cfg.name] - net_contact_forces = wp.to_torch(contact_sensor.data.net_forces_w_history) + net_contact_forces = contact_sensor.data.net_forces_w_history.torch # check if any contact force exceeds the threshold return torch.any( torch.max(torch.linalg.norm(net_contact_forces[:, :, sensor_cfg.body_ids], dim=-1), dim=1)[0] > threshold, dim=1 diff --git a/source/isaaclab/isaaclab/envs/ui/viewport_camera_controller.py b/source/isaaclab/isaaclab/envs/ui/viewport_camera_controller.py index 277982e7a2c9..509b17033397 100644 --- a/source/isaaclab/isaaclab/envs/ui/viewport_camera_controller.py +++ b/source/isaaclab/isaaclab/envs/ui/viewport_camera_controller.py @@ -12,7 +12,6 @@ import numpy as np import torch -import warp as wp from isaaclab.assets.articulation.articulation import Articulation @@ -162,7 +161,7 @@ def update_view_to_asset_root(self, asset_name: str): # set origin type to asset_root self.cfg.origin_type = "asset_root" # update the camera origins (convert Warp array to torch tensor first, then index) - root_pos = wp.to_torch(self._env.scene[self.cfg.asset_name].data.root_pos_w) + root_pos = self._env.scene[self.cfg.asset_name].data.root_pos_w.torch self.viewer_origin = root_pos[self.cfg.env_index] # update the camera view self.update_view_location() @@ -196,7 +195,7 @@ def update_view_to_asset_body(self, asset_name: str, body_name: str): # set origin type to asset_body self.cfg.origin_type = "asset_body" # update the camera origins (convert Warp array to torch tensor first, then index) - body_pos = wp.to_torch(self._env.scene[self.cfg.asset_name].data.body_pos_w) + body_pos = self._env.scene[self.cfg.asset_name].data.body_pos_w.torch self.viewer_origin = body_pos[self.cfg.env_index, body_id].squeeze(0) # update the camera view self.update_view_location() diff --git a/source/isaaclab/isaaclab/envs/utils/io_descriptors.py b/source/isaaclab/isaaclab/envs/utils/io_descriptors.py index ccf1feb43ec3..3a9d54f12e5a 100644 --- a/source/isaaclab/isaaclab/envs/utils/io_descriptors.py +++ b/source/isaaclab/isaaclab/envs/utils/io_descriptors.py @@ -20,8 +20,6 @@ import functools import inspect -import warp as wp - @configclass class GenericActionIODescriptor: @@ -321,7 +319,7 @@ def record_joint_pos_offsets(output: torch.Tensor, descriptor: GenericObservatio ids = kwargs["asset_cfg"].joint_ids # Get the offsets of the joints for the first robot in the scene. # This assumes that all robots have the same joint offsets. - descriptor.joint_pos_offsets = wp.to_torch(asset.data.default_joint_pos)[:, ids][0] + descriptor.joint_pos_offsets = asset.data.default_joint_pos.torch[:, ids][0] def record_joint_vel_offsets(output: torch.Tensor, descriptor: GenericObservationIODescriptor, **kwargs): @@ -338,7 +336,7 @@ def record_joint_vel_offsets(output: torch.Tensor, descriptor: GenericObservatio ids = kwargs["asset_cfg"].joint_ids # Get the offsets of the joints for the first robot in the scene. # This assumes that all robots have the same joint offsets. - descriptor.joint_vel_offsets = wp.to_torch(asset.data.default_joint_vel)[:, ids][0] + descriptor.joint_vel_offsets = asset.data.default_joint_vel.torch[:, ids][0] def export_articulations_data(env: ManagerBasedEnv) -> dict[str, dict[str, list[float]]]: @@ -358,25 +356,25 @@ def export_articulations_data(env: ManagerBasedEnv) -> dict[str, dict[str, list[ articulation_joint_data[articulation_name] = {} articulation_joint_data[articulation_name]["joint_names"] = articulation.joint_names articulation_joint_data[articulation_name]["default_joint_pos"] = ( - articulation.data.default_joint_pos[0].detach().cpu().numpy().tolist() + articulation.data.default_joint_pos.torch[0].detach().cpu().numpy().tolist() ) articulation_joint_data[articulation_name]["default_joint_vel"] = ( - articulation.data.default_joint_vel[0].detach().cpu().numpy().tolist() + articulation.data.default_joint_vel.torch[0].detach().cpu().numpy().tolist() ) articulation_joint_data[articulation_name]["default_joint_pos_limits"] = ( - articulation.data.default_joint_pos_limits[0].detach().cpu().numpy().tolist() + articulation.data.default_joint_pos_limits.torch[0].detach().cpu().numpy().tolist() ) articulation_joint_data[articulation_name]["default_joint_damping"] = ( - articulation.data.joint_damping[0].detach().cpu().numpy().tolist() + articulation.data.joint_damping.torch[0].detach().cpu().numpy().tolist() ) articulation_joint_data[articulation_name]["default_joint_stiffness"] = ( - articulation.data.joint_stiffness[0].detach().cpu().numpy().tolist() + articulation.data.joint_stiffness.torch[0].detach().cpu().numpy().tolist() ) articulation_joint_data[articulation_name]["default_joint_friction"] = ( - articulation.data.joint_friction_coeff[0].detach().cpu().numpy().tolist() + articulation.data.joint_friction_coeff.torch[0].detach().cpu().numpy().tolist() ) articulation_joint_data[articulation_name]["default_joint_armature"] = ( - articulation.data.joint_armature[0].detach().cpu().numpy().tolist() + articulation.data.joint_armature.torch[0].detach().cpu().numpy().tolist() ) return articulation_joint_data diff --git a/source/isaaclab/isaaclab/scene/interactive_scene.py b/source/isaaclab/isaaclab/scene/interactive_scene.py index da6f5eff75d7..096572bee6e1 100644 --- a/source/isaaclab/isaaclab/scene/interactive_scene.py +++ b/source/isaaclab/isaaclab/scene/interactive_scene.py @@ -612,30 +612,30 @@ def get_state(self, is_relative: bool = False) -> dict[str, dict[str, dict[str, state["articulation"] = dict() for asset_name, articulation in self._articulations.items(): asset_state = dict() - asset_state["root_pose"] = wp.to_torch(articulation.data.root_pose_w).clone() + asset_state["root_pose"] = articulation.data.root_pose_w.torch.clone() if is_relative: asset_state["root_pose"][:, :3] -= self.env_origins - asset_state["root_velocity"] = wp.to_torch(articulation.data.root_vel_w).clone() - asset_state["joint_position"] = wp.to_torch(articulation.data.joint_pos).clone() - asset_state["joint_velocity"] = wp.to_torch(articulation.data.joint_vel).clone() + asset_state["root_velocity"] = articulation.data.root_vel_w.torch.clone() + asset_state["joint_position"] = articulation.data.joint_pos.torch.clone() + asset_state["joint_velocity"] = articulation.data.joint_vel.torch.clone() state["articulation"][asset_name] = asset_state # deformable objects state["deformable_object"] = dict() for asset_name, deformable_object in self._deformable_objects.items(): asset_state = dict() - asset_state["nodal_position"] = wp.to_torch(deformable_object.data.nodal_pos_w).clone() + asset_state["nodal_position"] = deformable_object.data.nodal_pos_w.torch.clone() if is_relative: asset_state["nodal_position"][:, :3] -= self.env_origins - asset_state["nodal_velocity"] = wp.to_torch(deformable_object.data.nodal_vel_w).clone() + asset_state["nodal_velocity"] = deformable_object.data.nodal_vel_w.torch.clone() state["deformable_object"][asset_name] = asset_state # rigid objects state["rigid_object"] = dict() for asset_name, rigid_object in self._rigid_objects.items(): asset_state = dict() - asset_state["root_pose"] = wp.to_torch(rigid_object.data.root_pose_w).clone() + asset_state["root_pose"] = rigid_object.data.root_pose_w.torch.clone() if is_relative: asset_state["root_pose"][:, :3] -= self.env_origins - asset_state["root_velocity"] = wp.to_torch(rigid_object.data.root_vel_w).clone() + asset_state["root_velocity"] = rigid_object.data.root_vel_w.torch.clone() state["rigid_object"][asset_name] = asset_state # surface grippers state["gripper"] = dict() diff --git a/source/isaaclab/isaaclab/sensors/camera/camera.py b/source/isaaclab/isaaclab/sensors/camera/camera.py index 22c96af1779e..86ab80fd3c74 100644 --- a/source/isaaclab/isaaclab/sensors/camera/camera.py +++ b/source/isaaclab/isaaclab/sensors/camera/camera.py @@ -599,14 +599,14 @@ def _update_poses(self, env_ids: Sequence[int]): if len(self._sensor_prims) == 0: raise RuntimeError("Camera prim is None. Please call 'sim.play()' first.") - # get the poses from the view (returns wp.array, convert to torch) + # get the poses from the view (returns ProxyArray, use .torch for tensor access) if env_ids is not None and not isinstance(env_ids, torch.Tensor): env_ids = torch.tensor(env_ids, dtype=torch.int32, device=self._device) indices = wp.from_torch(env_ids.to(dtype=torch.int32), dtype=wp.int32) if env_ids is not None else None - pos_wp, quat_wp = self._view.get_world_poses(indices) - self._data.pos_w[env_ids] = wp.to_torch(pos_wp) + pos_w, quat_w = self._view.get_world_poses(indices) + self._data.pos_w[env_ids] = pos_w.torch self._data.quat_w_world[env_ids] = convert_camera_frame_orientation_convention( - wp.to_torch(quat_wp), origin="opengl", target="world" + quat_w.torch, origin="opengl", target="world" ) # notify renderer of updated poses (guarded in case called before initialization completes) if self._render_data is not None: diff --git a/source/isaaclab/isaaclab/sensors/contact_sensor/base_contact_sensor.py b/source/isaaclab/isaaclab/sensors/contact_sensor/base_contact_sensor.py index b0784b77d4dd..fb2b4a59b168 100644 --- a/source/isaaclab/isaaclab/sensors/contact_sensor/base_contact_sensor.py +++ b/source/isaaclab/isaaclab/sensors/contact_sensor/base_contact_sensor.py @@ -13,6 +13,7 @@ import warp as wp import isaaclab.utils.string as string_utils +from isaaclab.utils.warp import ProxyArray from ..sensor_base import SensorBase from .base_contact_sensor_data import BaseContactSensorData @@ -136,7 +137,7 @@ def find_sensors(self, name_keys: str | Sequence[str], preserve_order: bool = Fa return string_utils.resolve_matching_names(name_keys, self.body_names, preserve_order) @abstractmethod - def compute_first_contact(self, dt: float, abs_tol: float = 1.0e-8) -> wp.array: + def compute_first_contact(self, dt: float, abs_tol: float = 1.0e-8) -> ProxyArray: """Checks if bodies that have established contact within the last :attr:`dt` seconds. This function checks if the bodies have established contact within the last :attr:`dt` seconds @@ -164,7 +165,7 @@ def compute_first_contact(self, dt: float, abs_tol: float = 1.0e-8) -> wp.array: raise NotImplementedError(f"Compute first contact is not implemented for {self.__class__.__name__}.") @abstractmethod - def compute_first_air(self, dt: float, abs_tol: float = 1.0e-8) -> wp.array: + def compute_first_air(self, dt: float, abs_tol: float = 1.0e-8) -> ProxyArray: """Checks if bodies that have broken contact within the last :attr:`dt` seconds. This function checks if the bodies have broken contact within the last :attr:`dt` seconds diff --git a/source/isaaclab/isaaclab/sensors/contact_sensor/base_contact_sensor_data.py b/source/isaaclab/isaaclab/sensors/contact_sensor/base_contact_sensor_data.py index 74acf5092761..38702be28470 100644 --- a/source/isaaclab/isaaclab/sensors/contact_sensor/base_contact_sensor_data.py +++ b/source/isaaclab/isaaclab/sensors/contact_sensor/base_contact_sensor_data.py @@ -9,7 +9,7 @@ from abc import ABC, abstractmethod -import warp as wp +from isaaclab.utils.warp import ProxyArray class BaseContactSensorData(ABC): @@ -21,7 +21,7 @@ class BaseContactSensorData(ABC): @property @abstractmethod - def pose_w(self) -> wp.array | None: + def pose_w(self) -> ProxyArray | None: """Pose of the sensor origin in world frame. None if :attr:`ContactSensorCfg.track_pose` is False. @@ -30,7 +30,7 @@ def pose_w(self) -> wp.array | None: @property @abstractmethod - def pos_w(self) -> wp.array | None: + def pos_w(self) -> ProxyArray | None: """Position of the sensor origin in world frame. Shape is (num_instances, num_sensors), dtype = wp.vec3f. In torch this resolves to @@ -42,7 +42,7 @@ def pos_w(self) -> wp.array | None: @property @abstractmethod - def quat_w(self) -> wp.array | None: + def quat_w(self) -> ProxyArray | None: """Orientation of the sensor origin in world frame. Shape is (num_instances, num_sensors), dtype = wp.quatf. In torch this resolves to @@ -54,7 +54,7 @@ def quat_w(self) -> wp.array | None: @property @abstractmethod - def net_forces_w(self) -> wp.array | None: + def net_forces_w(self) -> ProxyArray | None: """The net normal contact forces in world frame. Shape is (num_instances, num_sensors), dtype = wp.vec3f. In torch this resolves to @@ -64,7 +64,7 @@ def net_forces_w(self) -> wp.array | None: @property @abstractmethod - def net_forces_w_history(self) -> wp.array | None: + def net_forces_w_history(self) -> ProxyArray | None: """History of net normal contact forces. Shape is (num_instances, history_length, num_sensors), dtype = wp.vec3f. In torch this resolves to @@ -74,7 +74,7 @@ def net_forces_w_history(self) -> wp.array | None: @property @abstractmethod - def force_matrix_w(self) -> wp.array | None: + def force_matrix_w(self) -> ProxyArray | None: """Normal contact forces filtered between sensor and filtered bodies. Shape is (num_instances, num_sensors, num_filter_shapes), dtype = wp.vec3f. In torch this resolves to @@ -86,7 +86,7 @@ def force_matrix_w(self) -> wp.array | None: @property @abstractmethod - def force_matrix_w_history(self) -> wp.array | None: + def force_matrix_w_history(self) -> ProxyArray | None: """History of filtered contact forces. Shape is (num_instances, history_length, num_sensors, num_filter_shapes), dtype = wp.vec3f. @@ -98,7 +98,7 @@ def force_matrix_w_history(self) -> wp.array | None: @property @abstractmethod - def contact_pos_w(self) -> wp.array | None: + def contact_pos_w(self) -> ProxyArray | None: """Average position of contact points. Shape is (num_instances, num_sensors, num_filter_shapes), dtype = wp.vec3f. In torch this resolves to @@ -110,7 +110,7 @@ def contact_pos_w(self) -> wp.array | None: @property @abstractmethod - def friction_forces_w(self) -> wp.array | None: + def friction_forces_w(self) -> ProxyArray | None: """Sum of friction forces. Shape is (num_instances, num_sensors, num_filter_shapes), dtype = wp.vec3f. In torch this resolves to @@ -122,7 +122,7 @@ def friction_forces_w(self) -> wp.array | None: @property @abstractmethod - def last_air_time(self) -> wp.array | None: + def last_air_time(self) -> ProxyArray | None: """Time spent in air before last contact. Shape is (num_instances, num_sensors), dtype = wp.float32. @@ -133,7 +133,7 @@ def last_air_time(self) -> wp.array | None: @property @abstractmethod - def current_air_time(self) -> wp.array | None: + def current_air_time(self) -> ProxyArray | None: """Time spent in air since last detach. Shape is (num_instances, num_sensors), dtype = wp.float32. @@ -144,7 +144,7 @@ def current_air_time(self) -> wp.array | None: @property @abstractmethod - def last_contact_time(self) -> wp.array | None: + def last_contact_time(self) -> ProxyArray | None: """Time spent in contact before last detach. Shape is (num_instances, num_sensors), dtype = wp.float32. @@ -155,7 +155,7 @@ def last_contact_time(self) -> wp.array | None: @property @abstractmethod - def current_contact_time(self) -> wp.array | None: + def current_contact_time(self) -> ProxyArray | None: """Time spent in contact since last contact. Shape is (num_instances, num_sensors), dtype = wp.float32. diff --git a/source/isaaclab/isaaclab/sensors/frame_transformer/base_frame_transformer_data.py b/source/isaaclab/isaaclab/sensors/frame_transformer/base_frame_transformer_data.py index 3b28a7b17d00..85386abfc37b 100644 --- a/source/isaaclab/isaaclab/sensors/frame_transformer/base_frame_transformer_data.py +++ b/source/isaaclab/isaaclab/sensors/frame_transformer/base_frame_transformer_data.py @@ -9,7 +9,7 @@ from abc import ABC, abstractmethod -import warp as wp +from isaaclab.utils.warp import ProxyArray class BaseFrameTransformerData(ABC): @@ -30,7 +30,7 @@ def target_frame_names(self) -> list[str]: @property @abstractmethod - def target_pose_source(self) -> wp.array | None: + def target_pose_source(self) -> ProxyArray | None: """Pose of the target frame(s) relative to source frame. Shape is (num_instances, num_target_frames), dtype = wp.transformf. In torch this resolves to @@ -40,7 +40,7 @@ def target_pose_source(self) -> wp.array | None: @property @abstractmethod - def target_pos_source(self) -> wp.array: + def target_pos_source(self) -> ProxyArray: """Position of the target frame(s) relative to source frame. Shape is (num_instances, num_target_frames), dtype = wp.vec3f. In torch this resolves to @@ -50,7 +50,7 @@ def target_pos_source(self) -> wp.array: @property @abstractmethod - def target_quat_source(self) -> wp.array: + def target_quat_source(self) -> ProxyArray: """Orientation of the target frame(s) relative to source frame. Shape is (num_instances, num_target_frames), dtype = wp.quatf. In torch this resolves to @@ -60,7 +60,7 @@ def target_quat_source(self) -> wp.array: @property @abstractmethod - def target_pose_w(self) -> wp.array | None: + def target_pose_w(self) -> ProxyArray | None: """Pose of the target frame(s) after offset in world frame. Shape is (num_instances, num_target_frames), dtype = wp.transformf. In torch this resolves to @@ -70,7 +70,7 @@ def target_pose_w(self) -> wp.array | None: @property @abstractmethod - def target_pos_w(self) -> wp.array: + def target_pos_w(self) -> ProxyArray: """Position of the target frame(s) after offset in world frame. Shape is (num_instances, num_target_frames), dtype = wp.vec3f. In torch this resolves to @@ -80,7 +80,7 @@ def target_pos_w(self) -> wp.array: @property @abstractmethod - def target_quat_w(self) -> wp.array: + def target_quat_w(self) -> ProxyArray: """Orientation of the target frame(s) after offset in world frame. Shape is (num_instances, num_target_frames), dtype = wp.quatf. In torch this resolves to @@ -90,7 +90,7 @@ def target_quat_w(self) -> wp.array: @property @abstractmethod - def source_pose_w(self) -> wp.array | None: + def source_pose_w(self) -> ProxyArray | None: """Pose of the source frame after offset in world frame. Shape is (num_instances,), dtype = wp.transformf. In torch this resolves to (num_instances, 7). @@ -100,7 +100,7 @@ def source_pose_w(self) -> wp.array | None: @property @abstractmethod - def source_pos_w(self) -> wp.array: + def source_pos_w(self) -> ProxyArray: """Position of the source frame after offset in world frame. Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). @@ -109,7 +109,7 @@ def source_pos_w(self) -> wp.array: @property @abstractmethod - def source_quat_w(self) -> wp.array: + def source_quat_w(self) -> ProxyArray: """Orientation of the source frame after offset in world frame. Shape is (num_instances,), dtype = wp.quatf. In torch this resolves to (num_instances, 4). diff --git a/source/isaaclab/isaaclab/sensors/imu/base_imu_data.py b/source/isaaclab/isaaclab/sensors/imu/base_imu_data.py index b1f175d4ebd4..bf8bf7bf2fc6 100644 --- a/source/isaaclab/isaaclab/sensors/imu/base_imu_data.py +++ b/source/isaaclab/isaaclab/sensors/imu/base_imu_data.py @@ -9,7 +9,7 @@ from abc import ABC, abstractmethod -import warp as wp +from isaaclab.utils.warp import ProxyArray class BaseImuData(ABC): @@ -24,7 +24,7 @@ class BaseImuData(ABC): @property @abstractmethod - def ang_vel_b(self) -> wp.array: + def ang_vel_b(self) -> ProxyArray: """IMU frame angular velocity relative to the world expressed in IMU frame [rad/s]. Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). @@ -33,7 +33,7 @@ def ang_vel_b(self) -> wp.array: @property @abstractmethod - def lin_acc_b(self) -> wp.array: + def lin_acc_b(self) -> ProxyArray: """Linear acceleration (proper) in the IMU frame [m/s^2]. Zero in freefall, +g upward at rest. diff --git a/source/isaaclab/isaaclab/sensors/pva/base_pva_data.py b/source/isaaclab/isaaclab/sensors/pva/base_pva_data.py index 17c56e3f3bb5..41dde4b11065 100644 --- a/source/isaaclab/isaaclab/sensors/pva/base_pva_data.py +++ b/source/isaaclab/isaaclab/sensors/pva/base_pva_data.py @@ -9,7 +9,7 @@ from abc import ABC, abstractmethod -import warp as wp +from isaaclab.utils.warp import ProxyArray class BasePvaData(ABC): @@ -21,7 +21,7 @@ class BasePvaData(ABC): @property @abstractmethod - def pose_w(self) -> wp.array | None: + def pose_w(self) -> ProxyArray | None: """Pose of the sensor origin in world frame [m, unitless]. Shape is (num_instances,), dtype = wp.transformf. In torch this resolves to (num_instances, 7). @@ -31,7 +31,7 @@ def pose_w(self) -> wp.array | None: @property @abstractmethod - def pos_w(self) -> wp.array: + def pos_w(self) -> ProxyArray: """Position of the sensor origin in world frame [m]. Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). @@ -40,7 +40,7 @@ def pos_w(self) -> wp.array: @property @abstractmethod - def quat_w(self) -> wp.array: + def quat_w(self) -> ProxyArray: """Orientation of the sensor origin in world frame. Shape is (num_instances,), dtype = wp.quatf. In torch this resolves to (num_instances, 4). @@ -50,7 +50,7 @@ def quat_w(self) -> wp.array: @property @abstractmethod - def projected_gravity_b(self) -> wp.array: + def projected_gravity_b(self) -> ProxyArray: """Gravity direction unit vector projected on the PVA frame. Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). @@ -59,7 +59,7 @@ def projected_gravity_b(self) -> wp.array: @property @abstractmethod - def lin_vel_b(self) -> wp.array: + def lin_vel_b(self) -> ProxyArray: """PVA frame linear velocity relative to the world expressed in PVA frame [m/s]. Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). @@ -68,7 +68,7 @@ def lin_vel_b(self) -> wp.array: @property @abstractmethod - def ang_vel_b(self) -> wp.array: + def ang_vel_b(self) -> ProxyArray: """PVA frame angular velocity relative to the world expressed in PVA frame [rad/s]. Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). @@ -77,7 +77,7 @@ def ang_vel_b(self) -> wp.array: @property @abstractmethod - def lin_acc_b(self) -> wp.array: + def lin_acc_b(self) -> ProxyArray: """Linear acceleration (coordinate) in the PVA frame [m/s^2]. Equal to -g in freefall, zero at rest. @@ -88,7 +88,7 @@ def lin_acc_b(self) -> wp.array: @property @abstractmethod - def ang_acc_b(self) -> wp.array: + def ang_acc_b(self) -> ProxyArray: """PVA frame angular acceleration relative to the world expressed in PVA frame [rad/s^2]. Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). diff --git a/source/isaaclab/isaaclab/sensors/ray_caster/multi_mesh_ray_caster.py b/source/isaaclab/isaaclab/sensors/ray_caster/multi_mesh_ray_caster.py index 9f3d692b13cd..f9d0493a0f61 100644 --- a/source/isaaclab/isaaclab/sensors/ray_caster/multi_mesh_ray_caster.py +++ b/source/isaaclab/isaaclab/sensors/ray_caster/multi_mesh_ray_caster.py @@ -350,8 +350,8 @@ def _update_mesh_transforms(self) -> None: continue # update position of the target meshes - pos_wp, quat_wp = view.get_world_poses(None) - pos_w, ori_w = wp.to_torch(pos_wp), wp.to_torch(quat_wp) + pos_w, ori_w = view.get_world_poses(None) + pos_w, ori_w = pos_w.torch, ori_w.torch pos_w = pos_w.squeeze(0) if len(pos_w.shape) == 3 else pos_w ori_w = ori_w.squeeze(0) if len(ori_w.shape) == 3 else ori_w diff --git a/source/isaaclab/isaaclab/sensors/ray_caster/multi_mesh_ray_caster_camera.py b/source/isaaclab/isaaclab/sensors/ray_caster/multi_mesh_ray_caster_camera.py index f184c28b20e2..a14c16f3d0d9 100644 --- a/source/isaaclab/isaaclab/sensors/ray_caster/multi_mesh_ray_caster_camera.py +++ b/source/isaaclab/isaaclab/sensors/ray_caster/multi_mesh_ray_caster_camera.py @@ -178,8 +178,8 @@ def _update_ray_infos(self, env_mask: wp.array): # Compute camera world poses by composing view pose with sensor offset indices = wp.from_torch(env_ids.to(dtype=torch.int32), dtype=wp.int32) - pos_wp, quat_wp = self._view.get_world_poses(indices) - pos_w, quat_w = wp.to_torch(pos_wp), wp.to_torch(quat_wp) + pos_w, quat_w = self._view.get_world_poses(indices) + pos_w, quat_w = pos_w.torch, quat_w.torch pos_w, quat_w = math_utils.combine_frame_transforms( pos_w, quat_w, self._offset_pos[env_ids], self._offset_quat[env_ids] ) diff --git a/source/isaaclab/isaaclab/sensors/ray_caster/ray_caster.py b/source/isaaclab/isaaclab/sensors/ray_caster/ray_caster.py index 1e23ee00c1b6..ab8ca8b7d4ce 100644 --- a/source/isaaclab/isaaclab/sensors/ray_caster/ray_caster.py +++ b/source/isaaclab/isaaclab/sensors/ray_caster/ray_caster.py @@ -264,9 +264,9 @@ def _get_view_transforms_wp(self) -> wp.array: ``(tx, ty, tz, qx, qy, qz, qw)`` per element, matching the quaternion convention returned by :class:`~isaaclab.sim.views.FrameView`. """ - pos_wp, quat_wp = self._view.get_world_poses() - pos_torch = wp.to_torch(pos_wp).reshape(-1, 3) - quat_torch = wp.to_torch(quat_wp).reshape(-1, 4) + pos_w, quat_w = self._view.get_world_poses() + pos_torch = pos_w.torch.reshape(-1, 3) + quat_torch = quat_w.torch.reshape(-1, 4) poses = torch.cat([pos_torch, quat_torch], dim=-1).contiguous() return wp.from_torch(poses).view(wp.transformf) diff --git a/source/isaaclab/isaaclab/sensors/ray_caster/ray_caster_camera.py b/source/isaaclab/isaaclab/sensors/ray_caster/ray_caster_camera.py index 17bb1e601980..257b25698deb 100644 --- a/source/isaaclab/isaaclab/sensors/ray_caster/ray_caster_camera.py +++ b/source/isaaclab/isaaclab/sensors/ray_caster/ray_caster_camera.py @@ -172,10 +172,9 @@ def reset(self, env_ids: Sequence[int] | None = None, env_mask: wp.array | None # reset the data # note: this recomputation is useful if one performs events such as randomizations on the camera poses. indices = wp.from_torch(env_ids.to(dtype=torch.int32), dtype=wp.int32) if env_ids is not None else None - pos_wp, quat_wp = self._view.get_world_poses(indices) - pos_w, quat_w = wp.to_torch(pos_wp).clone(), wp.to_torch(quat_wp).clone() + pos_w, quat_w = self._view.get_world_poses(indices) pos_w, quat_w = math_utils.combine_frame_transforms( - pos_w, quat_w, self._offset_pos[env_ids], self._offset_quat[env_ids] + pos_w.torch.clone(), quat_w.torch.clone(), self._offset_pos[env_ids], self._offset_quat[env_ids] ) self._data.pos_w[env_ids] = pos_w self._data.quat_w_world[env_ids] = quat_w @@ -218,27 +217,27 @@ def set_world_poses( # get current positions indices = wp.from_torch(env_ids.to(dtype=torch.int32), dtype=wp.int32) if env_ids is not None else None - pos_wp, quat_wp = self._view.get_world_poses(indices) - pos_w, quat_w = wp.to_torch(pos_wp), wp.to_torch(quat_wp) + pos_w, quat_w = self._view.get_world_poses(indices) + pos_w_torch = pos_w.torch + quat_w_torch = quat_w.torch if positions is not None: # transform to camera frame - pos_offset_world_frame = positions - pos_w - self._offset_pos[env_ids] = math_utils.quat_apply(math_utils.quat_inv(quat_w), pos_offset_world_frame) + pos_offset_world_frame = positions - pos_w_torch + self._offset_pos[env_ids] = math_utils.quat_apply(math_utils.quat_inv(quat_w_torch), pos_offset_world_frame) if orientations is not None: # convert rotation matrix from input convention to world quat_w_set = math_utils.convert_camera_frame_orientation_convention( orientations, origin=convention, target="world" ) - self._offset_quat[env_ids] = math_utils.quat_mul(math_utils.quat_inv(quat_w), quat_w_set) + self._offset_quat[env_ids] = math_utils.quat_mul(math_utils.quat_inv(quat_w_torch), quat_w_set) # update the data - pos_wp2, quat_wp2 = self._view.get_world_poses(indices) - pos_w, quat_w = wp.to_torch(pos_wp2).clone(), wp.to_torch(quat_wp2).clone() - pos_w, quat_w = math_utils.combine_frame_transforms( - pos_w, quat_w, self._offset_pos[env_ids], self._offset_quat[env_ids] + pos_w2, quat_w2 = self._view.get_world_poses(indices) + pos_w_out, quat_w_out = math_utils.combine_frame_transforms( + pos_w2.torch.clone(), quat_w2.torch.clone(), self._offset_pos[env_ids], self._offset_quat[env_ids] ) - self._data.pos_w[env_ids] = pos_w - self._data.quat_w_world[env_ids] = quat_w + self._data.pos_w[env_ids] = pos_w_out + self._data.quat_w_world[env_ids] = quat_w_out def set_world_poses_from_view( self, eyes: torch.Tensor, targets: torch.Tensor, env_ids: Sequence[int] | None = None @@ -581,7 +580,8 @@ def _compute_view_world_poses(self, env_ids: Sequence[int]) -> tuple[torch.Tenso .. deprecated v2.3.1: This function will be removed in a future release. Call - ``self._view.get_world_poses(indices)`` directly instead. + ``self._view.get_world_poses(indices)`` directly instead. The returned + ProxyArray pair exposes ``.warp`` and ``.torch`` accessors. Returns: A tuple of the position (in meters) and quaternion (x, y, z, w). @@ -594,8 +594,8 @@ def _compute_view_world_poses(self, env_ids: Sequence[int]) -> tuple[torch.Tenso ) indices = wp.from_torch(env_ids.to(dtype=torch.int32), dtype=wp.int32) if env_ids is not None else None - pos_wp, quat_wp = self._view.get_world_poses(indices) - return wp.to_torch(pos_wp).clone(), wp.to_torch(quat_wp).clone() + pos_w, quat_w = self._view.get_world_poses(indices) + return pos_w.torch.clone(), quat_w.torch.clone() def _compute_camera_world_poses(self, env_ids: Sequence[int]) -> tuple[torch.Tensor, torch.Tensor]: """Computes the pose of the camera in the world frame. @@ -608,8 +608,9 @@ def _compute_camera_world_poses(self, env_ids: Sequence[int]) -> tuple[torch.Ten .. code-block:: python indices = wp.from_torch(env_ids.to(dtype=torch.int32), dtype=wp.int32) - pos_wp, quat_wp = self._view.get_world_poses(indices) - pos_w, quat_w = wp.to_torch(pos_wp).clone(), wp.to_torch(quat_wp).clone() + pos_w, quat_w = self._view.get_world_poses(indices) + # The returned ProxyArray pair exposes .warp and .torch accessors + pos_w, quat_w = pos_w.torch.clone(), quat_w.torch.clone() pos_w, quat_w = math_utils.combine_frame_transforms( pos_w, quat_w, self._offset_pos[env_ids], self._offset_quat[env_ids] ) @@ -623,6 +624,7 @@ def _compute_camera_world_poses(self, env_ids: Sequence[int]) -> tuple[torch.Ten ) indices = wp.from_torch(env_ids.to(dtype=torch.int32), dtype=wp.int32) if env_ids is not None else None - pos_wp, quat_wp = self._view.get_world_poses(indices) - pos_w, quat_w = wp.to_torch(pos_wp).clone(), wp.to_torch(quat_wp).clone() - return math_utils.combine_frame_transforms(pos_w, quat_w, self._offset_pos[env_ids], self._offset_quat[env_ids]) + pos_w, quat_w = self._view.get_world_poses(indices) + return math_utils.combine_frame_transforms( + pos_w.torch.clone(), quat_w.torch.clone(), self._offset_pos[env_ids], self._offset_quat[env_ids] + ) diff --git a/source/isaaclab/isaaclab/sensors/ray_caster/ray_caster_data.py b/source/isaaclab/isaaclab/sensors/ray_caster/ray_caster_data.py index c317c96f78b2..6ccd1f5729f1 100644 --- a/source/isaaclab/isaaclab/sensors/ray_caster/ray_caster_data.py +++ b/source/isaaclab/isaaclab/sensors/ray_caster/ray_caster_data.py @@ -7,13 +7,15 @@ import warp as wp +from isaaclab.utils.warp import ProxyArray + class RayCasterData: """Data container for the ray-cast sensor. - All public properties return :class:`wp.array` objects backed by device memory. - Use :func:`wp.to_torch` at the call-site when a PyTorch tensor is needed, e.g. - ``wp.to_torch(sensor.data.ray_hits_w)``. + Public properties return :class:`~isaaclab.utils.warp.ProxyArray` wrappers. + Use ``.torch`` for a cached zero-copy :class:`torch.Tensor` view or + ``.warp`` for the underlying :class:`warp.array`. """ def __init__(self): @@ -21,45 +23,42 @@ def __init__(self): self._quat_w: wp.array | None = None self._ray_hits_w: wp.array | None = None - # Zero-copy torch views; kept alive to prevent GC of the underlying warp buffers. - # Not surfaced as public API — callers should use wp.to_torch() at the call-site. - self._pos_w_torch = None - self._quat_w_torch = None - self._ray_hits_w_torch = None + # _pos_w_ta / _quat_w_ta / _ray_hits_w_ta are created in create_buffers(). + # Accessing the public properties before create_buffers() raises AttributeError. @property - def pos_w(self) -> wp.array | None: + def pos_w(self) -> ProxyArray: """Position of the sensor origin in world frame [m]. Shape is (N,), dtype ``wp.vec3f``. In torch this resolves to (N, 3), - where N is the number of sensors. Use :func:`wp.to_torch` to obtain a - :class:`torch.Tensor` view without copying data. + where N is the number of sensors. Use ``.warp`` for the underlying + ``wp.array`` or ``.torch`` for a cached zero-copy ``torch.Tensor`` view. """ - return self._pos_w + return self._pos_w_ta @property - def quat_w(self) -> wp.array | None: + def quat_w(self) -> ProxyArray: """Orientation of the sensor origin in quaternion (x, y, z, w) in world frame. Shape is (N,), dtype ``wp.quatf``. In torch this resolves to (N, 4), - where N is the number of sensors. Use :func:`wp.to_torch` to obtain a - :class:`torch.Tensor` view without copying data. + where N is the number of sensors. Use ``.warp`` for the underlying + ``wp.array`` or ``.torch`` for a cached zero-copy ``torch.Tensor`` view. """ - return self._quat_w + return self._quat_w_ta @property - def ray_hits_w(self) -> wp.array | None: + def ray_hits_w(self) -> ProxyArray: """The ray hit positions in the world frame [m]. Shape is (N, B), dtype ``wp.vec3f``. In torch this resolves to (N, B, 3), where N is the number of sensors and B is the number of rays per sensor. - Contains ``inf`` for missed hits. Use :func:`wp.to_torch` to obtain a - :class:`torch.Tensor` view without copying data. + Contains ``inf`` for missed hits. Use ``.warp`` for the underlying + ``wp.array`` or ``.torch`` for a cached zero-copy ``torch.Tensor`` view. """ - return self._ray_hits_w + return self._ray_hits_w_ta def create_buffers(self, num_envs: int, num_rays: int, device: str) -> None: - """Create internal warp buffers and corresponding zero-copy torch views. + """Create internal warp buffers and their :class:`ProxyArray` wrappers. Args: num_envs: Number of environments / sensors. @@ -72,6 +71,6 @@ def create_buffers(self, num_envs: int, num_rays: int, device: str) -> None: self._quat_w = wp.zeros(num_envs, dtype=wp.quatf, device=device) self._ray_hits_w = wp.zeros((num_envs, num_rays), dtype=wp.vec3f, device=device) - self._pos_w_torch = wp.to_torch(self._pos_w) - self._quat_w_torch = wp.to_torch(self._quat_w) - self._ray_hits_w_torch = wp.to_torch(self._ray_hits_w) + self._pos_w_ta = ProxyArray(self._pos_w) + self._quat_w_ta = ProxyArray(self._quat_w) + self._ray_hits_w_ta = ProxyArray(self._ray_hits_w) diff --git a/source/isaaclab/isaaclab/sim/views/base_frame_view.py b/source/isaaclab/isaaclab/sim/views/base_frame_view.py index fc59c2ed83ab..d58cb0fbb534 100644 --- a/source/isaaclab/isaaclab/sim/views/base_frame_view.py +++ b/source/isaaclab/isaaclab/sim/views/base_frame_view.py @@ -11,6 +11,8 @@ import warp as wp +from isaaclab.utils.warp import ProxyArray + class BaseFrameView(abc.ABC): """Abstract interface for reading and writing world-space transforms of multiple prims. @@ -20,7 +22,7 @@ class BaseFrameView(abc.ABC): :class:`~isaaclab.sim.views.FrameView` selects the correct implementation at runtime based on the active physics backend. - All getters return ``wp.array``. Setters accept ``wp.array``. + All pose getters return :class:`~isaaclab.utils.warp.ProxyArray`. Setters accept ``wp.array``. """ @property @@ -29,15 +31,23 @@ def count(self) -> int: """Number of prims in this view.""" ... + @property + @abc.abstractmethod + def device(self) -> str: + """Device where arrays are allocated (``"cpu"`` or ``"cuda:0"``).""" + ... + @abc.abstractmethod - def get_world_poses(self, indices: wp.array | None = None) -> tuple[wp.array, wp.array]: + def get_world_poses(self, indices: wp.array | None = None) -> tuple[ProxyArray, ProxyArray]: """Get world-space positions and orientations for prims in the view. Args: indices: Subset of prims to query. ``None`` means all prims. Returns: - A tuple ``(positions (M, 3), orientations (M, 4))`` as ``wp.array``. + A tuple ``(positions, orientations)`` of :class:`~isaaclab.utils.warp.ProxyArray` + wrappers. Use ``.warp`` for the underlying ``wp.array`` or ``.torch`` for a + cached zero-copy ``torch.Tensor`` view. """ ... @@ -58,14 +68,16 @@ def set_world_poses( ... @abc.abstractmethod - def get_local_poses(self, indices: wp.array | None = None) -> tuple[wp.array, wp.array]: + def get_local_poses(self, indices: wp.array | None = None) -> tuple[ProxyArray, ProxyArray]: """Get local-space positions and orientations for prims in the view. Args: indices: Subset of prims to query. ``None`` means all prims. Returns: - A tuple ``(translations (M, 3), orientations (M, 4))`` as ``wp.array``. + A tuple ``(translations, orientations)`` of :class:`~isaaclab.utils.warp.ProxyArray` + wrappers. Use ``.warp`` for the underlying ``wp.array`` or ``.torch`` for a + cached zero-copy ``torch.Tensor`` view. """ ... diff --git a/source/isaaclab/isaaclab/sim/views/usd_frame_view.py b/source/isaaclab/isaaclab/sim/views/usd_frame_view.py index 4421fa5391ea..7730c3dd735d 100644 --- a/source/isaaclab/isaaclab/sim/views/usd_frame_view.py +++ b/source/isaaclab/isaaclab/sim/views/usd_frame_view.py @@ -14,6 +14,7 @@ from pxr import Gf, Sdf, Usd, UsdGeom, Vt import isaaclab.sim as sim_utils +from isaaclab.utils.warp import ProxyArray from .base_frame_view import BaseFrameView @@ -34,7 +35,7 @@ class UsdFrameView(BaseFrameView): For GPU-accelerated Fabric operations, use the PhysX backend variant obtained via :class:`~isaaclab.sim.views.FrameView`. - All getters return ``wp.array``. Setters accept ``wp.array``. + Pose getters return :class:`~isaaclab.utils.warp.ProxyArray`. Setters accept ``wp.array``. .. note:: **Transform Requirements:** @@ -252,14 +253,16 @@ def set_visibility(self, visibility: torch.Tensor, indices: wp.array | None = No # Getters # ------------------------------------------------------------------ - def get_world_poses(self, indices: wp.array | None = None) -> tuple[wp.array, wp.array]: + def get_world_poses(self, indices: wp.array | None = None) -> tuple[ProxyArray, ProxyArray]: """Get world-space poses for prims in the view. Args: indices: Indices of prims to get poses for. Defaults to None (all prims). Returns: - A tuple of ``(positions, orientations)`` as ``wp.array``. + A tuple ``(positions, orientations)`` of :class:`~isaaclab.utils.warp.ProxyArray` + wrappers. Use ``.warp`` for the underlying ``wp.array`` or ``.torch`` for a + cached zero-copy ``torch.Tensor`` view. """ indices_list = self._resolve_indices(indices) @@ -274,19 +277,20 @@ def get_world_poses(self, indices: wp.array | None = None) -> tuple[wp.array, wp positions[idx] = prim_tf.ExtractTranslation() orientations[idx] = prim_tf.ExtractRotationQuat() - return ( - wp.array(np.array(positions, dtype=np.float32), dtype=wp.float32, device=self._device), - wp.array(np.array(orientations, dtype=np.float32), dtype=wp.float32, device=self._device), - ) + pos_wp = wp.array(np.array(positions, dtype=np.float32), dtype=wp.float32, device=self._device) + quat_wp = wp.array(np.array(orientations, dtype=np.float32), dtype=wp.float32, device=self._device) + return ProxyArray(pos_wp), ProxyArray(quat_wp) - def get_local_poses(self, indices: wp.array | None = None) -> tuple[wp.array, wp.array]: + def get_local_poses(self, indices: wp.array | None = None) -> tuple[ProxyArray, ProxyArray]: """Get local-space poses for prims in the view. Args: indices: Indices of prims to get poses for. Defaults to None (all prims). Returns: - A tuple of ``(translations, orientations)`` as ``wp.array``. + A tuple ``(translations, orientations)`` of :class:`~isaaclab.utils.warp.ProxyArray` + wrappers. Use ``.warp`` for the underlying ``wp.array`` or ``.torch`` for a + cached zero-copy ``torch.Tensor`` view. """ indices_list = self._resolve_indices(indices) @@ -301,10 +305,9 @@ def get_local_poses(self, indices: wp.array | None = None) -> tuple[wp.array, wp translations[idx] = prim_tf.ExtractTranslation() orientations[idx] = prim_tf.ExtractRotationQuat() - return ( - wp.array(np.array(translations, dtype=np.float32), dtype=wp.float32, device=self._device), - wp.array(np.array(orientations, dtype=np.float32), dtype=wp.float32, device=self._device), - ) + pos_wp = wp.array(np.array(translations, dtype=np.float32), dtype=wp.float32, device=self._device) + quat_wp = wp.array(np.array(orientations, dtype=np.float32), dtype=wp.float32, device=self._device) + return ProxyArray(pos_wp), ProxyArray(quat_wp) def get_scales(self, indices: wp.array | None = None) -> wp.array: """Get scales for prims in the view. diff --git a/source/isaaclab/isaaclab/test/mock_interfaces/assets/mock_articulation.py b/source/isaaclab/isaaclab/test/mock_interfaces/assets/mock_articulation.py index ddbee003152c..aa72bf2bbfd9 100644 --- a/source/isaaclab/isaaclab/test/mock_interfaces/assets/mock_articulation.py +++ b/source/isaaclab/isaaclab/test/mock_interfaces/assets/mock_articulation.py @@ -14,6 +14,8 @@ import torch import warp as wp +from isaaclab.utils.warp import ProxyArray + try: from isaaclab.assets.articulation.base_articulation_data import BaseArticulationData except (ImportError, ModuleNotFoundError): @@ -148,6 +150,55 @@ def __init__( self._spatial_tendon_limit_stiffness: wp.array | None = None self._spatial_tendon_offset: wp.array | None = None + # ProxyArray caches + self._default_root_pose_ta: ProxyArray | None = None + self._default_root_vel_ta: ProxyArray | None = None + self._default_joint_pos_ta: ProxyArray | None = None + self._default_joint_vel_ta: ProxyArray | None = None + self._joint_pos_target_ta: ProxyArray | None = None + self._joint_vel_target_ta: ProxyArray | None = None + self._joint_effort_target_ta: ProxyArray | None = None + self._computed_torque_ta: ProxyArray | None = None + self._applied_torque_ta: ProxyArray | None = None + self._joint_stiffness_ta: ProxyArray | None = None + self._joint_damping_ta: ProxyArray | None = None + self._joint_armature_ta: ProxyArray | None = None + self._joint_friction_coeff_ta: ProxyArray | None = None + self._joint_dynamic_friction_coeff_ta: ProxyArray | None = None + self._joint_viscous_friction_coeff_ta: ProxyArray | None = None + self._joint_pos_limits_ta: ProxyArray | None = None + self._joint_vel_limits_ta: ProxyArray | None = None + self._joint_effort_limits_ta: ProxyArray | None = None + self._soft_joint_pos_limits_ta: ProxyArray | None = None + self._soft_joint_vel_limits_ta: ProxyArray | None = None + self._gear_ratio_ta: ProxyArray | None = None + self._joint_pos_ta: ProxyArray | None = None + self._joint_vel_ta: ProxyArray | None = None + self._joint_acc_ta: ProxyArray | None = None + self._root_link_pose_w_ta: ProxyArray | None = None + self._root_link_vel_w_ta: ProxyArray | None = None + self._root_com_pose_w_ta: ProxyArray | None = None + self._root_com_vel_w_ta: ProxyArray | None = None + self._body_link_pose_w_ta: ProxyArray | None = None + self._body_link_vel_w_ta: ProxyArray | None = None + self._body_com_pose_w_ta: ProxyArray | None = None + self._body_com_vel_w_ta: ProxyArray | None = None + self._body_com_acc_w_ta: ProxyArray | None = None + self._body_com_pose_b_ta: ProxyArray | None = None + self._body_mass_ta: ProxyArray | None = None + self._body_inertia_ta: ProxyArray | None = None + self._body_incoming_joint_wrench_b_ta: ProxyArray | None = None + self._fixed_tendon_stiffness_ta: ProxyArray | None = None + self._fixed_tendon_damping_ta: ProxyArray | None = None + self._fixed_tendon_limit_stiffness_ta: ProxyArray | None = None + self._fixed_tendon_rest_length_ta: ProxyArray | None = None + self._fixed_tendon_offset_ta: ProxyArray | None = None + self._fixed_tendon_pos_limits_ta: ProxyArray | None = None + self._spatial_tendon_stiffness_ta: ProxyArray | None = None + self._spatial_tendon_damping_ta: ProxyArray | None = None + self._spatial_tendon_limit_stiffness_ta: ProxyArray | None = None + self._spatial_tendon_offset_ta: ProxyArray | None = None + # -- Helper for identity quaternion -- def _identity_quat(self, *shape: int) -> wp.array: """Create identity quaternion array (w, x, y, z) = (1, 0, 0, 0).""" @@ -158,515 +209,672 @@ def _identity_quat(self, *shape: int) -> wp.array: # -- Default state properties -- @property - def default_root_pose(self) -> wp.array: + def default_root_pose(self) -> ProxyArray: """Default root pose. dtype=wp.transformf, shape: (N,).""" if self._default_root_pose is None: pose_np = np.zeros((self._num_instances, 7), dtype=np.float32) pose_np[..., 6] = 1.0 # identity quat qw=1, transformf layout: (px,py,pz,qx,qy,qz,qw) self._default_root_pose = wp.array(pose_np, dtype=wp.float32, device=self.device).view(wp.transformf) - return self._default_root_pose + self._default_root_pose_ta = None + if self._default_root_pose_ta is None: + self._default_root_pose_ta = ProxyArray(self._default_root_pose) + return self._default_root_pose_ta @property - def default_root_vel(self) -> wp.array: + def default_root_vel(self) -> ProxyArray: """Default root velocity. dtype=wp.spatial_vectorf, shape: (N,).""" if self._default_root_vel is None: self._default_root_vel = wp.zeros((self._num_instances, 6), dtype=wp.float32, device=self.device).view( wp.spatial_vectorf ) - return self._default_root_vel + self._default_root_vel_ta = None + if self._default_root_vel_ta is None: + self._default_root_vel_ta = ProxyArray(self._default_root_vel) + return self._default_root_vel_ta @property - def default_root_state(self) -> wp.array: + def default_root_state(self) -> ProxyArray: """Default root state [pose(7), vel(6)]. Shape: (N, 13).""" - pose_t = wp.to_torch(self.default_root_pose) - vel_t = wp.to_torch(self.default_root_vel) - return wp.from_torch(torch.cat([pose_t, vel_t], dim=-1)) + pose_t = self.default_root_pose.torch + vel_t = self.default_root_vel.torch + return ProxyArray(wp.from_torch(torch.cat([pose_t, vel_t], dim=-1))) @property - def default_joint_pos(self) -> wp.array: + def default_joint_pos(self) -> ProxyArray: """Default joint positions. Shape: (N, num_joints).""" if self._default_joint_pos is None: - return wp.zeros((self._num_instances, self._num_joints), dtype=wp.float32, device=self.device) - return self._default_joint_pos + self._default_joint_pos = wp.zeros( + (self._num_instances, self._num_joints), dtype=wp.float32, device=self.device + ) + self._default_joint_pos_ta = None + if self._default_joint_pos_ta is None: + self._default_joint_pos_ta = ProxyArray(self._default_joint_pos) + return self._default_joint_pos_ta @property - def default_joint_vel(self) -> wp.array: + def default_joint_vel(self) -> ProxyArray: """Default joint velocities. Shape: (N, num_joints).""" if self._default_joint_vel is None: - return wp.zeros((self._num_instances, self._num_joints), dtype=wp.float32, device=self.device) - return self._default_joint_vel + self._default_joint_vel = wp.zeros( + (self._num_instances, self._num_joints), dtype=wp.float32, device=self.device + ) + self._default_joint_vel_ta = None + if self._default_joint_vel_ta is None: + self._default_joint_vel_ta = ProxyArray(self._default_joint_vel) + return self._default_joint_vel_ta # -- Joint command properties -- @property - def joint_pos_target(self) -> wp.array: + def joint_pos_target(self) -> ProxyArray: """Joint position targets. Shape: (N, num_joints).""" if self._joint_pos_target is None: - return wp.zeros((self._num_instances, self._num_joints), dtype=wp.float32, device=self.device) - return self._joint_pos_target + self._joint_pos_target = wp.zeros( + (self._num_instances, self._num_joints), dtype=wp.float32, device=self.device + ) + self._joint_pos_target_ta = None + if self._joint_pos_target_ta is None: + self._joint_pos_target_ta = ProxyArray(self._joint_pos_target) + return self._joint_pos_target_ta @property - def joint_vel_target(self) -> wp.array: + def joint_vel_target(self) -> ProxyArray: """Joint velocity targets. Shape: (N, num_joints).""" if self._joint_vel_target is None: - return wp.zeros((self._num_instances, self._num_joints), dtype=wp.float32, device=self.device) - return self._joint_vel_target + self._joint_vel_target = wp.zeros( + (self._num_instances, self._num_joints), dtype=wp.float32, device=self.device + ) + self._joint_vel_target_ta = None + if self._joint_vel_target_ta is None: + self._joint_vel_target_ta = ProxyArray(self._joint_vel_target) + return self._joint_vel_target_ta @property - def joint_effort_target(self) -> wp.array: + def joint_effort_target(self) -> ProxyArray: """Joint effort targets. Shape: (N, num_joints).""" if self._joint_effort_target is None: - return wp.zeros((self._num_instances, self._num_joints), dtype=wp.float32, device=self.device) - return self._joint_effort_target + self._joint_effort_target = wp.zeros( + (self._num_instances, self._num_joints), dtype=wp.float32, device=self.device + ) + self._joint_effort_target_ta = None + if self._joint_effort_target_ta is None: + self._joint_effort_target_ta = ProxyArray(self._joint_effort_target) + return self._joint_effort_target_ta @property - def computed_torque(self) -> wp.array: + def computed_torque(self) -> ProxyArray: """Computed torques before clipping. Shape: (N, num_joints).""" if self._computed_torque is None: - return wp.zeros((self._num_instances, self._num_joints), dtype=wp.float32, device=self.device) - return self._computed_torque + self._computed_torque = wp.zeros( + (self._num_instances, self._num_joints), dtype=wp.float32, device=self.device + ) + self._computed_torque_ta = None + if self._computed_torque_ta is None: + self._computed_torque_ta = ProxyArray(self._computed_torque) + return self._computed_torque_ta @property - def applied_torque(self) -> wp.array: + def applied_torque(self) -> ProxyArray: """Applied torques after clipping. Shape: (N, num_joints).""" if self._applied_torque is None: - return wp.zeros((self._num_instances, self._num_joints), dtype=wp.float32, device=self.device) - return self._applied_torque + self._applied_torque = wp.zeros( + (self._num_instances, self._num_joints), dtype=wp.float32, device=self.device + ) + self._applied_torque_ta = None + if self._applied_torque_ta is None: + self._applied_torque_ta = ProxyArray(self._applied_torque) + return self._applied_torque_ta # -- Joint properties -- @property - def joint_stiffness(self) -> wp.array: + def joint_stiffness(self) -> ProxyArray: """Joint stiffness. Shape: (N, num_joints).""" if self._joint_stiffness is None: - return wp.zeros((self._num_instances, self._num_joints), dtype=wp.float32, device=self.device) - return self._joint_stiffness + self._joint_stiffness = wp.zeros( + (self._num_instances, self._num_joints), dtype=wp.float32, device=self.device + ) + self._joint_stiffness_ta = None + if self._joint_stiffness_ta is None: + self._joint_stiffness_ta = ProxyArray(self._joint_stiffness) + return self._joint_stiffness_ta @property - def joint_damping(self) -> wp.array: + def joint_damping(self) -> ProxyArray: """Joint damping. Shape: (N, num_joints).""" if self._joint_damping is None: - return wp.zeros((self._num_instances, self._num_joints), dtype=wp.float32, device=self.device) - return self._joint_damping + self._joint_damping = wp.zeros( + (self._num_instances, self._num_joints), dtype=wp.float32, device=self.device + ) + self._joint_damping_ta = None + if self._joint_damping_ta is None: + self._joint_damping_ta = ProxyArray(self._joint_damping) + return self._joint_damping_ta @property - def joint_armature(self) -> wp.array: + def joint_armature(self) -> ProxyArray: """Joint armature. Shape: (N, num_joints).""" if self._joint_armature is None: - return wp.zeros((self._num_instances, self._num_joints), dtype=wp.float32, device=self.device) - return self._joint_armature + self._joint_armature = wp.zeros( + (self._num_instances, self._num_joints), dtype=wp.float32, device=self.device + ) + self._joint_armature_ta = None + if self._joint_armature_ta is None: + self._joint_armature_ta = ProxyArray(self._joint_armature) + return self._joint_armature_ta @property - def joint_friction_coeff(self) -> wp.array: + def joint_friction_coeff(self) -> ProxyArray: """Joint static friction coefficient. Shape: (N, num_joints).""" if self._joint_friction_coeff is None: - return wp.zeros((self._num_instances, self._num_joints), dtype=wp.float32, device=self.device) - return self._joint_friction_coeff + self._joint_friction_coeff = wp.zeros( + (self._num_instances, self._num_joints), dtype=wp.float32, device=self.device + ) + self._joint_friction_coeff_ta = None + if self._joint_friction_coeff_ta is None: + self._joint_friction_coeff_ta = ProxyArray(self._joint_friction_coeff) + return self._joint_friction_coeff_ta @property - def joint_dynamic_friction_coeff(self) -> wp.array: + def joint_dynamic_friction_coeff(self) -> ProxyArray: """Joint dynamic friction coefficient. Shape: (N, num_joints).""" if self._joint_dynamic_friction_coeff is None: - return wp.zeros((self._num_instances, self._num_joints), dtype=wp.float32, device=self.device) - return self._joint_dynamic_friction_coeff + self._joint_dynamic_friction_coeff = wp.zeros( + (self._num_instances, self._num_joints), dtype=wp.float32, device=self.device + ) + self._joint_dynamic_friction_coeff_ta = None + if self._joint_dynamic_friction_coeff_ta is None: + self._joint_dynamic_friction_coeff_ta = ProxyArray(self._joint_dynamic_friction_coeff) + return self._joint_dynamic_friction_coeff_ta @property - def joint_viscous_friction_coeff(self) -> wp.array: + def joint_viscous_friction_coeff(self) -> ProxyArray: """Joint viscous friction coefficient. Shape: (N, num_joints).""" if self._joint_viscous_friction_coeff is None: - return wp.zeros((self._num_instances, self._num_joints), dtype=wp.float32, device=self.device) - return self._joint_viscous_friction_coeff + self._joint_viscous_friction_coeff = wp.zeros( + (self._num_instances, self._num_joints), dtype=wp.float32, device=self.device + ) + self._joint_viscous_friction_coeff_ta = None + if self._joint_viscous_friction_coeff_ta is None: + self._joint_viscous_friction_coeff_ta = ProxyArray(self._joint_viscous_friction_coeff) + return self._joint_viscous_friction_coeff_ta @property - def joint_pos_limits(self) -> wp.array: + def joint_pos_limits(self) -> ProxyArray: """Joint position limits [lower, upper]. dtype=wp.vec2f, shape: (N, num_joints).""" if self._joint_pos_limits is None: np_limits = np.zeros((self._num_instances, self._num_joints, 2), dtype=np.float32) np_limits[..., 0] = -float("inf") np_limits[..., 1] = float("inf") - return wp.array(np_limits, dtype=wp.float32, device=self.device).view(wp.vec2f) - return self._joint_pos_limits + self._joint_pos_limits = wp.array(np_limits, dtype=wp.float32, device=self.device).view(wp.vec2f) + self._joint_pos_limits_ta = None + if self._joint_pos_limits_ta is None: + self._joint_pos_limits_ta = ProxyArray(self._joint_pos_limits) + return self._joint_pos_limits_ta @property - def joint_vel_limits(self) -> wp.array: + def joint_vel_limits(self) -> ProxyArray: """Joint velocity limits. Shape: (N, num_joints).""" if self._joint_vel_limits is None: - return wp.full( + self._joint_vel_limits = wp.full( (self._num_instances, self._num_joints), dtype=wp.float32, value=float("inf"), device=self.device ) - return self._joint_vel_limits + self._joint_vel_limits_ta = None + if self._joint_vel_limits_ta is None: + self._joint_vel_limits_ta = ProxyArray(self._joint_vel_limits) + return self._joint_vel_limits_ta @property - def joint_effort_limits(self) -> wp.array: + def joint_effort_limits(self) -> ProxyArray: """Joint effort limits. Shape: (N, num_joints).""" if self._joint_effort_limits is None: - return wp.full( + self._joint_effort_limits = wp.full( (self._num_instances, self._num_joints), dtype=wp.float32, value=float("inf"), device=self.device ) - return self._joint_effort_limits + self._joint_effort_limits_ta = None + if self._joint_effort_limits_ta is None: + self._joint_effort_limits_ta = ProxyArray(self._joint_effort_limits) + return self._joint_effort_limits_ta @property - def soft_joint_pos_limits(self) -> wp.array: + def soft_joint_pos_limits(self) -> ProxyArray: """Soft joint position limits. Shape: (N, num_joints, 2).""" if self._soft_joint_pos_limits is None: - return wp.clone(self.joint_pos_limits, self.device) - return self._soft_joint_pos_limits + self._soft_joint_pos_limits = wp.clone(self.joint_pos_limits.warp, self.device) + self._soft_joint_pos_limits_ta = None + if self._soft_joint_pos_limits_ta is None: + self._soft_joint_pos_limits_ta = ProxyArray(self._soft_joint_pos_limits) + return self._soft_joint_pos_limits_ta @property - def soft_joint_vel_limits(self) -> wp.array: + def soft_joint_vel_limits(self) -> ProxyArray: """Soft joint velocity limits. Shape: (N, num_joints).""" if self._soft_joint_vel_limits is None: - return wp.clone(self.joint_vel_limits, self.device) - return self._soft_joint_vel_limits + self._soft_joint_vel_limits = wp.clone(self.joint_vel_limits.warp, self.device) + self._soft_joint_vel_limits_ta = None + if self._soft_joint_vel_limits_ta is None: + self._soft_joint_vel_limits_ta = ProxyArray(self._soft_joint_vel_limits) + return self._soft_joint_vel_limits_ta @property - def gear_ratio(self) -> wp.array: + def gear_ratio(self) -> ProxyArray: """Gear ratio. Shape: (N, num_joints).""" if self._gear_ratio is None: - return wp.ones((self._num_instances, self._num_joints), dtype=wp.float32, device=self.device) - return self._gear_ratio + self._gear_ratio = wp.ones((self._num_instances, self._num_joints), dtype=wp.float32, device=self.device) + self._gear_ratio_ta = None + if self._gear_ratio_ta is None: + self._gear_ratio_ta = ProxyArray(self._gear_ratio) + return self._gear_ratio_ta # -- Joint state properties -- @property - def joint_pos(self) -> wp.array: + def joint_pos(self) -> ProxyArray: """Joint positions. Shape: (N, num_joints).""" if self._joint_pos is None: - return wp.zeros((self._num_instances, self._num_joints), dtype=wp.float32, device=self.device) - return self._joint_pos + self._joint_pos = wp.zeros((self._num_instances, self._num_joints), dtype=wp.float32, device=self.device) + self._joint_pos_ta = None + if self._joint_pos_ta is None: + self._joint_pos_ta = ProxyArray(self._joint_pos) + return self._joint_pos_ta @property - def joint_vel(self) -> wp.array: + def joint_vel(self) -> ProxyArray: """Joint velocities. Shape: (N, num_joints).""" if self._joint_vel is None: - return wp.zeros((self._num_instances, self._num_joints), dtype=wp.float32, device=self.device) - return self._joint_vel + self._joint_vel = wp.zeros((self._num_instances, self._num_joints), dtype=wp.float32, device=self.device) + self._joint_vel_ta = None + if self._joint_vel_ta is None: + self._joint_vel_ta = ProxyArray(self._joint_vel) + return self._joint_vel_ta @property - def joint_acc(self) -> wp.array: + def joint_acc(self) -> ProxyArray: """Joint accelerations. Shape: (N, num_joints).""" if self._joint_acc is None: - return wp.zeros((self._num_instances, self._num_joints), dtype=wp.float32, device=self.device) - return self._joint_acc + self._joint_acc = wp.zeros((self._num_instances, self._num_joints), dtype=wp.float32, device=self.device) + self._joint_acc_ta = None + if self._joint_acc_ta is None: + self._joint_acc_ta = ProxyArray(self._joint_acc) + return self._joint_acc_ta # -- Root state properties (link frame) -- @property - def root_link_pose_w(self) -> wp.array: + def root_link_pose_w(self) -> ProxyArray: """Root link pose in world frame. dtype=wp.transformf, shape: (N,).""" if self._root_link_pose_w is None: pose_np = np.zeros((self._num_instances, 7), dtype=np.float32) pose_np[..., 6] = 1.0 # identity quat qw=1, transformf layout: (px,py,pz,qx,qy,qz,qw) self._root_link_pose_w = wp.array(pose_np, dtype=wp.float32, device=self.device).view(wp.transformf) - return self._root_link_pose_w + self._root_link_pose_w_ta = None + if self._root_link_pose_w_ta is None: + self._root_link_pose_w_ta = ProxyArray(self._root_link_pose_w) + return self._root_link_pose_w_ta @property - def root_link_vel_w(self) -> wp.array: + def root_link_vel_w(self) -> ProxyArray: """Root link velocity in world frame. dtype=wp.spatial_vectorf, shape: (N,).""" if self._root_link_vel_w is None: self._root_link_vel_w = wp.zeros((self._num_instances, 6), dtype=wp.float32, device=self.device).view( wp.spatial_vectorf ) - return self._root_link_vel_w + self._root_link_vel_w_ta = None + if self._root_link_vel_w_ta is None: + self._root_link_vel_w_ta = ProxyArray(self._root_link_vel_w) + return self._root_link_vel_w_ta @property - def root_link_state_w(self) -> wp.array: + def root_link_state_w(self) -> ProxyArray: """Root link state in world frame. Shape: (N, 13).""" - pose_t = wp.to_torch(self.root_link_pose_w) - vel_t = wp.to_torch(self.root_link_vel_w) - return wp.from_torch(torch.cat([pose_t, vel_t], dim=-1)) + pose_t = self.root_link_pose_w.torch + vel_t = self.root_link_vel_w.torch + return ProxyArray(wp.from_torch(torch.cat([pose_t, vel_t], dim=-1).contiguous())) # Sliced properties (zero-copy pointer arithmetic on transformf) @property - def root_link_pos_w(self) -> wp.array: + def root_link_pos_w(self) -> ProxyArray: """Root link position in world frame. Shape: (N,), dtype=wp.vec3f.""" - t = self.root_link_pose_w - return wp.array(ptr=t.ptr, shape=t.shape, dtype=wp.vec3f, strides=t.strides, device=self.device) + t = self.root_link_pose_w.warp + return ProxyArray(wp.array(ptr=t.ptr, shape=t.shape, dtype=wp.vec3f, strides=t.strides, device=self.device)) @property - def root_link_quat_w(self) -> wp.array: + def root_link_quat_w(self) -> ProxyArray: """Root link orientation in world frame (x, y, z, w). Shape: (N,), dtype=wp.quatf.""" - t = self.root_link_pose_w - return wp.array(ptr=t.ptr + 3 * 4, shape=t.shape, dtype=wp.quatf, strides=t.strides, device=self.device) + t = self.root_link_pose_w.warp + return ProxyArray( + wp.array(ptr=t.ptr + 3 * 4, shape=t.shape, dtype=wp.quatf, strides=t.strides, device=self.device) + ) @property - def root_link_lin_vel_w(self) -> wp.array: + def root_link_lin_vel_w(self) -> ProxyArray: """Root link linear velocity in world frame. Shape: (N,), dtype=wp.vec3f.""" - v = self.root_link_vel_w - return wp.array(ptr=v.ptr, shape=v.shape, dtype=wp.vec3f, strides=v.strides, device=self.device) + v = self.root_link_vel_w.warp + return ProxyArray(wp.array(ptr=v.ptr, shape=v.shape, dtype=wp.vec3f, strides=v.strides, device=self.device)) @property - def root_link_ang_vel_w(self) -> wp.array: + def root_link_ang_vel_w(self) -> ProxyArray: """Root link angular velocity in world frame. Shape: (N,), dtype=wp.vec3f.""" - v = self.root_link_vel_w - return wp.array(ptr=v.ptr + 3 * 4, shape=v.shape, dtype=wp.vec3f, strides=v.strides, device=self.device) + v = self.root_link_vel_w.warp + return ProxyArray( + wp.array(ptr=v.ptr + 3 * 4, shape=v.shape, dtype=wp.vec3f, strides=v.strides, device=self.device) + ) # -- Root state properties (CoM frame) -- @property - def root_com_pose_w(self) -> wp.array: + def root_com_pose_w(self) -> ProxyArray: """Root CoM pose in world frame. dtype=wp.transformf, shape: (N,).""" if self._root_com_pose_w is None: - self._root_com_pose_w = wp.clone(self.root_link_pose_w, self.device) - return self._root_com_pose_w + self._root_com_pose_w = wp.clone(self.root_link_pose_w.warp, self.device) + self._root_com_pose_w_ta = None + if self._root_com_pose_w_ta is None: + self._root_com_pose_w_ta = ProxyArray(self._root_com_pose_w) + return self._root_com_pose_w_ta @property - def root_com_vel_w(self) -> wp.array: + def root_com_vel_w(self) -> ProxyArray: """Root CoM velocity in world frame. dtype=wp.spatial_vectorf, shape: (N,).""" if self._root_com_vel_w is None: - self._root_com_vel_w = wp.clone(self.root_link_vel_w, self.device) - return self._root_com_vel_w + self._root_com_vel_w = wp.clone(self.root_link_vel_w.warp, self.device) + self._root_com_vel_w_ta = None + if self._root_com_vel_w_ta is None: + self._root_com_vel_w_ta = ProxyArray(self._root_com_vel_w) + return self._root_com_vel_w_ta @property - def root_com_state_w(self) -> wp.array: + def root_com_state_w(self) -> ProxyArray: """Root CoM state in world frame. Shape: (N, 13).""" - pose_t = wp.to_torch(self.root_com_pose_w) - vel_t = wp.to_torch(self.root_com_vel_w) - return wp.from_torch(torch.cat([pose_t, vel_t], dim=-1)) + pose_t = self.root_com_pose_w.torch + vel_t = self.root_com_vel_w.torch + return ProxyArray(wp.from_torch(torch.cat([pose_t, vel_t], dim=-1).contiguous())) @property - def root_state_w(self) -> wp.array: + def root_state_w(self) -> ProxyArray: """Root state (link pose + CoM velocity). Shape: (N, 13).""" - pose_t = wp.to_torch(self.root_link_pose_w) - vel_t = wp.to_torch(self.root_com_vel_w) - return wp.from_torch(torch.cat([pose_t, vel_t], dim=-1)) + pose_t = self.root_link_pose_w.torch + vel_t = self.root_com_vel_w.torch + return ProxyArray(wp.from_torch(torch.cat([pose_t, vel_t], dim=-1).contiguous())) # Sliced properties (zero-copy pointer arithmetic on transformf) @property - def root_com_pos_w(self) -> wp.array: + def root_com_pos_w(self) -> ProxyArray: """Root CoM position in world frame. Shape: (N,), dtype=wp.vec3f.""" - t = self.root_com_pose_w - return wp.array(ptr=t.ptr, shape=t.shape, dtype=wp.vec3f, strides=t.strides, device=self.device) + t = self.root_com_pose_w.warp + return ProxyArray(wp.array(ptr=t.ptr, shape=t.shape, dtype=wp.vec3f, strides=t.strides, device=self.device)) @property - def root_com_quat_w(self) -> wp.array: + def root_com_quat_w(self) -> ProxyArray: """Root CoM orientation in world frame. Shape: (N,), dtype=wp.quatf.""" - t = self.root_com_pose_w - return wp.array(ptr=t.ptr + 3 * 4, shape=t.shape, dtype=wp.quatf, strides=t.strides, device=self.device) + t = self.root_com_pose_w.warp + return ProxyArray( + wp.array(ptr=t.ptr + 3 * 4, shape=t.shape, dtype=wp.quatf, strides=t.strides, device=self.device) + ) @property - def root_com_lin_vel_w(self) -> wp.array: + def root_com_lin_vel_w(self) -> ProxyArray: """Root CoM linear velocity in world frame. Shape: (N,), dtype=wp.vec3f.""" - v = self.root_com_vel_w - return wp.array(ptr=v.ptr, shape=v.shape, dtype=wp.vec3f, strides=v.strides, device=self.device) + v = self.root_com_vel_w.warp + return ProxyArray(wp.array(ptr=v.ptr, shape=v.shape, dtype=wp.vec3f, strides=v.strides, device=self.device)) @property - def root_com_ang_vel_w(self) -> wp.array: + def root_com_ang_vel_w(self) -> ProxyArray: """Root CoM angular velocity in world frame. Shape: (N,), dtype=wp.vec3f.""" - v = self.root_com_vel_w - return wp.array(ptr=v.ptr + 3 * 4, shape=v.shape, dtype=wp.vec3f, strides=v.strides, device=self.device) + v = self.root_com_vel_w.warp + return ProxyArray( + wp.array(ptr=v.ptr + 3 * 4, shape=v.shape, dtype=wp.vec3f, strides=v.strides, device=self.device) + ) # -- Body state properties (link frame) -- @property - def body_link_pose_w(self) -> wp.array: + def body_link_pose_w(self) -> ProxyArray: """Body link poses in world frame. dtype=wp.transformf, shape: (N, num_bodies).""" if self._body_link_pose_w is None: pose_np = np.zeros((self._num_instances, self._num_bodies, 7), dtype=np.float32) pose_np[..., 6] = 1.0 # identity quat qw=1, transformf layout: (px,py,pz,qx,qy,qz,qw) self._body_link_pose_w = wp.array(pose_np, dtype=wp.float32, device=self.device).view(wp.transformf) - return self._body_link_pose_w + self._body_link_pose_w_ta = None + if self._body_link_pose_w_ta is None: + self._body_link_pose_w_ta = ProxyArray(self._body_link_pose_w) + return self._body_link_pose_w_ta @property - def body_link_vel_w(self) -> wp.array: + def body_link_vel_w(self) -> ProxyArray: """Body link velocities in world frame. dtype=wp.spatial_vectorf, shape: (N, num_bodies).""" if self._body_link_vel_w is None: self._body_link_vel_w = wp.zeros( (self._num_instances, self._num_bodies, 6), dtype=wp.float32, device=self.device ).view(wp.spatial_vectorf) - return self._body_link_vel_w + self._body_link_vel_w_ta = None + if self._body_link_vel_w_ta is None: + self._body_link_vel_w_ta = ProxyArray(self._body_link_vel_w) + return self._body_link_vel_w_ta @property - def body_link_state_w(self) -> wp.array: + def body_link_state_w(self) -> ProxyArray: """Body link states in world frame. Shape: (N, num_bodies, 13).""" - pose_t = wp.to_torch(self.body_link_pose_w) - vel_t = wp.to_torch(self.body_link_vel_w) - return wp.from_torch(torch.cat([pose_t, vel_t], dim=-1)) + pose_t = self.body_link_pose_w.torch + vel_t = self.body_link_vel_w.torch + return ProxyArray(wp.from_torch(torch.cat([pose_t, vel_t], dim=-1).contiguous())) # Sliced properties (zero-copy pointer arithmetic on transformf) @property - def body_link_pos_w(self) -> wp.array: + def body_link_pos_w(self) -> ProxyArray: """Body link positions. Shape: (N, num_bodies), dtype=wp.vec3f.""" - t = self.body_link_pose_w - return wp.array(ptr=t.ptr, shape=t.shape, dtype=wp.vec3f, strides=t.strides, device=self.device) + t = self.body_link_pose_w.warp + return ProxyArray(wp.array(ptr=t.ptr, shape=t.shape, dtype=wp.vec3f, strides=t.strides, device=self.device)) @property - def body_link_quat_w(self) -> wp.array: + def body_link_quat_w(self) -> ProxyArray: """Body link orientations. Shape: (N, num_bodies), dtype=wp.quatf.""" - t = self.body_link_pose_w - return wp.array(ptr=t.ptr + 3 * 4, shape=t.shape, dtype=wp.quatf, strides=t.strides, device=self.device) + t = self.body_link_pose_w.warp + return ProxyArray( + wp.array(ptr=t.ptr + 3 * 4, shape=t.shape, dtype=wp.quatf, strides=t.strides, device=self.device) + ) @property - def body_link_lin_vel_w(self) -> wp.array: + def body_link_lin_vel_w(self) -> ProxyArray: """Body link linear velocities in world frame. Shape: (N, num_bodies), dtype=wp.vec3f.""" - v = self.body_link_vel_w - return wp.array(ptr=v.ptr, shape=v.shape, dtype=wp.vec3f, strides=v.strides, device=self.device) + v = self.body_link_vel_w.warp + return ProxyArray(wp.array(ptr=v.ptr, shape=v.shape, dtype=wp.vec3f, strides=v.strides, device=self.device)) @property - def body_link_ang_vel_w(self) -> wp.array: + def body_link_ang_vel_w(self) -> ProxyArray: """Body link angular velocities in world frame. Shape: (N, num_bodies), dtype=wp.vec3f.""" - v = self.body_link_vel_w - return wp.array(ptr=v.ptr + 3 * 4, shape=v.shape, dtype=wp.vec3f, strides=v.strides, device=self.device) + v = self.body_link_vel_w.warp + return ProxyArray( + wp.array(ptr=v.ptr + 3 * 4, shape=v.shape, dtype=wp.vec3f, strides=v.strides, device=self.device) + ) # -- Body state properties (CoM frame) -- @property - def body_com_pose_w(self) -> wp.array: + def body_com_pose_w(self) -> ProxyArray: """Body CoM poses in world frame. dtype=wp.transformf, shape: (N, num_bodies).""" if self._body_com_pose_w is None: - self._body_com_pose_w = wp.clone(self.body_link_pose_w, self.device) - return self._body_com_pose_w + self._body_com_pose_w = wp.clone(self.body_link_pose_w.warp, self.device) + self._body_com_pose_w_ta = None + if self._body_com_pose_w_ta is None: + self._body_com_pose_w_ta = ProxyArray(self._body_com_pose_w) + return self._body_com_pose_w_ta @property - def body_com_vel_w(self) -> wp.array: + def body_com_vel_w(self) -> ProxyArray: """Body CoM velocities in world frame. dtype=wp.spatial_vectorf, shape: (N, num_bodies).""" if self._body_com_vel_w is None: - self._body_com_vel_w = wp.clone(self.body_link_vel_w, self.device) - return self._body_com_vel_w + self._body_com_vel_w = wp.clone(self.body_link_vel_w.warp, self.device) + self._body_com_vel_w_ta = None + if self._body_com_vel_w_ta is None: + self._body_com_vel_w_ta = ProxyArray(self._body_com_vel_w) + return self._body_com_vel_w_ta @property - def body_com_state_w(self) -> wp.array: + def body_com_state_w(self) -> ProxyArray: """Body CoM states in world frame. Shape: (N, num_bodies, 13).""" - pose_t = wp.to_torch(self.body_com_pose_w) - vel_t = wp.to_torch(self.body_com_vel_w) - return wp.from_torch(torch.cat([pose_t, vel_t], dim=-1)) + pose_t = self.body_com_pose_w.torch + vel_t = self.body_com_vel_w.torch + return ProxyArray(wp.from_torch(torch.cat([pose_t, vel_t], dim=-1).contiguous())) @property - def body_state_w(self) -> wp.array: + def body_state_w(self) -> ProxyArray: """Body states (link pose + CoM velocity). Shape: (N, num_bodies, 13).""" - pose_t = wp.to_torch(self.body_link_pose_w) - vel_t = wp.to_torch(self.body_com_vel_w) - return wp.from_torch(torch.cat([pose_t, vel_t], dim=-1)) + pose_t = self.body_link_pose_w.torch + vel_t = self.body_com_vel_w.torch + return ProxyArray(wp.from_torch(torch.cat([pose_t, vel_t], dim=-1).contiguous())) @property - def body_com_acc_w(self) -> wp.array: + def body_com_acc_w(self) -> ProxyArray: """Body CoM accelerations in world frame. dtype=wp.spatial_vectorf, shape: (N, num_bodies).""" if self._body_com_acc_w is None: self._body_com_acc_w = wp.zeros( (self._num_instances, self._num_bodies, 6), dtype=wp.float32, device=self.device ).view(wp.spatial_vectorf) - return self._body_com_acc_w + self._body_com_acc_w_ta = None + if self._body_com_acc_w_ta is None: + self._body_com_acc_w_ta = ProxyArray(self._body_com_acc_w) + return self._body_com_acc_w_ta @property - def body_com_pose_b(self) -> wp.array: + def body_com_pose_b(self) -> ProxyArray: """Body CoM poses in body frame. dtype=wp.transformf, shape: (N, num_bodies).""" if self._body_com_pose_b is None: pose_np = np.zeros((self._num_instances, self._num_bodies, 7), dtype=np.float32) pose_np[..., 6] = 1.0 # identity quat qw=1, transformf layout: (px,py,pz,qx,qy,qz,qw) self._body_com_pose_b = wp.array(pose_np, dtype=wp.float32, device=self.device).view(wp.transformf) - return self._body_com_pose_b + self._body_com_pose_b_ta = None + if self._body_com_pose_b_ta is None: + self._body_com_pose_b_ta = ProxyArray(self._body_com_pose_b) + return self._body_com_pose_b_ta # Sliced properties (zero-copy pointer arithmetic on transformf) @property - def body_com_pos_w(self) -> wp.array: + def body_com_pos_w(self) -> ProxyArray: """Body CoM positions. Shape: (N, num_bodies), dtype=wp.vec3f.""" - t = self.body_com_pose_w - return wp.array(ptr=t.ptr, shape=t.shape, dtype=wp.vec3f, strides=t.strides, device=self.device) + t = self.body_com_pose_w.warp + return ProxyArray(wp.array(ptr=t.ptr, shape=t.shape, dtype=wp.vec3f, strides=t.strides, device=self.device)) @property - def body_com_quat_w(self) -> wp.array: + def body_com_quat_w(self) -> ProxyArray: """Body CoM orientations. Shape: (N, num_bodies), dtype=wp.quatf.""" - t = self.body_com_pose_w - return wp.array(ptr=t.ptr + 3 * 4, shape=t.shape, dtype=wp.quatf, strides=t.strides, device=self.device) + t = self.body_com_pose_w.warp + return ProxyArray( + wp.array(ptr=t.ptr + 3 * 4, shape=t.shape, dtype=wp.quatf, strides=t.strides, device=self.device) + ) @property - def body_com_lin_vel_w(self) -> wp.array: + def body_com_lin_vel_w(self) -> ProxyArray: """Body CoM linear velocities in world frame. Shape: (N, num_bodies), dtype=wp.vec3f.""" - v = self.body_com_vel_w - return wp.array(ptr=v.ptr, shape=v.shape, dtype=wp.vec3f, strides=v.strides, device=self.device) + v = self.body_com_vel_w.warp + return ProxyArray(wp.array(ptr=v.ptr, shape=v.shape, dtype=wp.vec3f, strides=v.strides, device=self.device)) @property - def body_com_ang_vel_w(self) -> wp.array: + def body_com_ang_vel_w(self) -> ProxyArray: """Body CoM angular velocities in world frame. Shape: (N, num_bodies), dtype=wp.vec3f.""" - v = self.body_com_vel_w - return wp.array(ptr=v.ptr + 3 * 4, shape=v.shape, dtype=wp.vec3f, strides=v.strides, device=self.device) + v = self.body_com_vel_w.warp + return ProxyArray( + wp.array(ptr=v.ptr + 3 * 4, shape=v.shape, dtype=wp.vec3f, strides=v.strides, device=self.device) + ) @property - def body_com_lin_acc_w(self) -> wp.array: + def body_com_lin_acc_w(self) -> ProxyArray: """Body CoM linear accelerations in world frame. Shape: (N, num_bodies), dtype=wp.vec3f.""" - a = self.body_com_acc_w - return wp.array(ptr=a.ptr, shape=a.shape, dtype=wp.vec3f, strides=a.strides, device=self.device) + a = self.body_com_acc_w.warp + return ProxyArray(wp.array(ptr=a.ptr, shape=a.shape, dtype=wp.vec3f, strides=a.strides, device=self.device)) @property - def body_com_ang_acc_w(self) -> wp.array: + def body_com_ang_acc_w(self) -> ProxyArray: """Body CoM angular accelerations in world frame. Shape: (N, num_bodies), dtype=wp.vec3f.""" - a = self.body_com_acc_w - return wp.array(ptr=a.ptr + 3 * 4, shape=a.shape, dtype=wp.vec3f, strides=a.strides, device=self.device) + a = self.body_com_acc_w.warp + return ProxyArray( + wp.array(ptr=a.ptr + 3 * 4, shape=a.shape, dtype=wp.vec3f, strides=a.strides, device=self.device) + ) @property - def body_com_pos_b(self) -> wp.array: + def body_com_pos_b(self) -> ProxyArray: """Body CoM positions in body frame. Shape: (N, num_bodies), dtype=wp.vec3f.""" - t = self.body_com_pose_b - return wp.array(ptr=t.ptr, shape=t.shape, dtype=wp.vec3f, strides=t.strides, device=self.device) + t = self.body_com_pose_b.warp + return ProxyArray(wp.array(ptr=t.ptr, shape=t.shape, dtype=wp.vec3f, strides=t.strides, device=self.device)) @property - def body_com_quat_b(self) -> wp.array: + def body_com_quat_b(self) -> ProxyArray: """Body CoM orientations in body frame. Shape: (N, num_bodies), dtype=wp.quatf.""" - t = self.body_com_pose_b - return wp.array(ptr=t.ptr + 3 * 4, shape=t.shape, dtype=wp.quatf, strides=t.strides, device=self.device) + t = self.body_com_pose_b.warp + return ProxyArray( + wp.array(ptr=t.ptr + 3 * 4, shape=t.shape, dtype=wp.quatf, strides=t.strides, device=self.device) + ) # -- Body properties -- @property - def body_mass(self) -> wp.array: + def body_mass(self) -> ProxyArray: """Body masses. Shape: (N, num_bodies).""" if self._body_mass is None: - return wp.ones((self._num_instances, self._num_bodies), dtype=wp.float32, device=self.device) - return self._body_mass + self._body_mass = wp.ones((self._num_instances, self._num_bodies), dtype=wp.float32, device=self.device) + self._body_mass_ta = None + if self._body_mass_ta is None: + self._body_mass_ta = ProxyArray(self._body_mass) + return self._body_mass_ta @property - def body_inertia(self) -> wp.array: + def body_inertia(self) -> ProxyArray: """Body inertias (flattened 3x3). Shape: (N, num_bodies, 9).""" if self._body_inertia is None: np_inertia = np.zeros((self._num_instances, self._num_bodies, 9), dtype=np.float32) np_inertia[..., 0] = 1.0 # Ixx np_inertia[..., 4] = 1.0 # Iyy np_inertia[..., 8] = 1.0 # Izz - return wp.array(np_inertia, dtype=wp.float32, device=self.device) - return self._body_inertia + self._body_inertia = wp.array(np_inertia, dtype=wp.float32, device=self.device) + self._body_inertia_ta = None + if self._body_inertia_ta is None: + self._body_inertia_ta = ProxyArray(self._body_inertia) + return self._body_inertia_ta @property - def body_incoming_joint_wrench_b(self) -> wp.array: + def body_incoming_joint_wrench_b(self) -> ProxyArray: """Body incoming joint wrenches. dtype=wp.spatial_vectorf, shape: (N, num_bodies).""" if self._body_incoming_joint_wrench_b is None: - return wp.zeros((self._num_instances, self._num_bodies, 6), dtype=wp.float32, device=self.device).view( - wp.spatial_vectorf - ) - return self._body_incoming_joint_wrench_b + self._body_incoming_joint_wrench_b = wp.zeros( + (self._num_instances, self._num_bodies, 6), dtype=wp.float32, device=self.device + ).view(wp.spatial_vectorf) + self._body_incoming_joint_wrench_b_ta = None + if self._body_incoming_joint_wrench_b_ta is None: + self._body_incoming_joint_wrench_b_ta = ProxyArray(self._body_incoming_joint_wrench_b) + return self._body_incoming_joint_wrench_b_ta # -- Derived properties -- @property - def projected_gravity_b(self) -> wp.array: + def projected_gravity_b(self) -> ProxyArray: """Gravity projection on base. Shape: (N,), dtype=wp.vec3f.""" np_gravity = np.zeros((self._num_instances, 3), dtype=np.float32) np_gravity[:, 2] = -1.0 - return wp.array(np_gravity, dtype=wp.float32, device=self.device).view(wp.vec3f) + return ProxyArray(wp.array(np_gravity, dtype=wp.float32, device=self.device).view(wp.vec3f)) @property - def heading_w(self) -> wp.array: + def heading_w(self) -> ProxyArray: """Yaw heading in world frame. Shape: (N,).""" - return wp.zeros((self._num_instances,), dtype=wp.float32, device=self.device) + return ProxyArray(wp.zeros((self._num_instances,), dtype=wp.float32, device=self.device)) @property - def root_link_lin_vel_b(self) -> wp.array: + def root_link_lin_vel_b(self) -> ProxyArray: """Root link linear velocity in body frame. Shape: (N,), dtype=wp.vec3f.""" - return wp.clone(self.root_link_lin_vel_w, self.device) + return ProxyArray(wp.clone(self.root_link_lin_vel_w.warp, self.device)) @property - def root_link_ang_vel_b(self) -> wp.array: + def root_link_ang_vel_b(self) -> ProxyArray: """Root link angular velocity in body frame. Shape: (N,), dtype=wp.vec3f.""" - return wp.clone(self.root_link_ang_vel_w, self.device) + return ProxyArray(wp.clone(self.root_link_ang_vel_w.warp, self.device)) @property - def root_com_lin_vel_b(self) -> wp.array: + def root_com_lin_vel_b(self) -> ProxyArray: """Root CoM linear velocity in body frame. Shape: (N,), dtype=wp.vec3f.""" - return wp.clone(self.root_com_lin_vel_w, self.device) + return ProxyArray(wp.clone(self.root_com_lin_vel_w.warp, self.device)) @property - def root_com_ang_vel_b(self) -> wp.array: + def root_com_ang_vel_b(self) -> ProxyArray: """Root CoM angular velocity in body frame. Shape: (N,), dtype=wp.vec3f.""" - return wp.clone(self.root_com_ang_vel_w, self.device) + return ProxyArray(wp.clone(self.root_com_ang_vel_w.warp, self.device)) # com_pos_b and com_quat_b are inherited from BaseArticulationData # (they delegate to body_com_pos_b and body_com_quat_b respectively) @@ -674,99 +882,147 @@ def root_com_ang_vel_b(self) -> wp.array: # -- Fixed tendon properties -- @property - def fixed_tendon_stiffness(self) -> wp.array: + def fixed_tendon_stiffness(self) -> ProxyArray: """Fixed tendon stiffness. Shape: (N, num_fixed_tendons).""" if self._num_fixed_tendons == 0: - return wp.zeros((self._num_instances, 0), dtype=wp.float32, device=self.device) + return ProxyArray(wp.zeros((self._num_instances, 0), dtype=wp.float32, device=self.device)) if self._fixed_tendon_stiffness is None: - return wp.zeros((self._num_instances, self._num_fixed_tendons), dtype=wp.float32, device=self.device) - return self._fixed_tendon_stiffness + self._fixed_tendon_stiffness = wp.zeros( + (self._num_instances, self._num_fixed_tendons), dtype=wp.float32, device=self.device + ) + self._fixed_tendon_stiffness_ta = None + if self._fixed_tendon_stiffness_ta is None: + self._fixed_tendon_stiffness_ta = ProxyArray(self._fixed_tendon_stiffness) + return self._fixed_tendon_stiffness_ta @property - def fixed_tendon_damping(self) -> wp.array: + def fixed_tendon_damping(self) -> ProxyArray: """Fixed tendon damping. Shape: (N, num_fixed_tendons).""" if self._num_fixed_tendons == 0: - return wp.zeros((self._num_instances, 0), dtype=wp.float32, device=self.device) + return ProxyArray(wp.zeros((self._num_instances, 0), dtype=wp.float32, device=self.device)) if self._fixed_tendon_damping is None: - return wp.zeros((self._num_instances, self._num_fixed_tendons), dtype=wp.float32, device=self.device) - return self._fixed_tendon_damping + self._fixed_tendon_damping = wp.zeros( + (self._num_instances, self._num_fixed_tendons), dtype=wp.float32, device=self.device + ) + self._fixed_tendon_damping_ta = None + if self._fixed_tendon_damping_ta is None: + self._fixed_tendon_damping_ta = ProxyArray(self._fixed_tendon_damping) + return self._fixed_tendon_damping_ta @property - def fixed_tendon_limit_stiffness(self) -> wp.array: + def fixed_tendon_limit_stiffness(self) -> ProxyArray: """Fixed tendon limit stiffness. Shape: (N, num_fixed_tendons).""" if self._num_fixed_tendons == 0: - return wp.zeros((self._num_instances, 0), dtype=wp.float32, device=self.device) + return ProxyArray(wp.zeros((self._num_instances, 0), dtype=wp.float32, device=self.device)) if self._fixed_tendon_limit_stiffness is None: - return wp.zeros((self._num_instances, self._num_fixed_tendons), dtype=wp.float32, device=self.device) - return self._fixed_tendon_limit_stiffness + self._fixed_tendon_limit_stiffness = wp.zeros( + (self._num_instances, self._num_fixed_tendons), dtype=wp.float32, device=self.device + ) + self._fixed_tendon_limit_stiffness_ta = None + if self._fixed_tendon_limit_stiffness_ta is None: + self._fixed_tendon_limit_stiffness_ta = ProxyArray(self._fixed_tendon_limit_stiffness) + return self._fixed_tendon_limit_stiffness_ta @property - def fixed_tendon_rest_length(self) -> wp.array: + def fixed_tendon_rest_length(self) -> ProxyArray: """Fixed tendon rest length. Shape: (N, num_fixed_tendons).""" if self._num_fixed_tendons == 0: - return wp.zeros((self._num_instances, 0), dtype=wp.float32, device=self.device) + return ProxyArray(wp.zeros((self._num_instances, 0), dtype=wp.float32, device=self.device)) if self._fixed_tendon_rest_length is None: - return wp.zeros((self._num_instances, self._num_fixed_tendons), dtype=wp.float32, device=self.device) - return self._fixed_tendon_rest_length + self._fixed_tendon_rest_length = wp.zeros( + (self._num_instances, self._num_fixed_tendons), dtype=wp.float32, device=self.device + ) + self._fixed_tendon_rest_length_ta = None + if self._fixed_tendon_rest_length_ta is None: + self._fixed_tendon_rest_length_ta = ProxyArray(self._fixed_tendon_rest_length) + return self._fixed_tendon_rest_length_ta @property - def fixed_tendon_offset(self) -> wp.array: + def fixed_tendon_offset(self) -> ProxyArray: """Fixed tendon offset. Shape: (N, num_fixed_tendons).""" if self._num_fixed_tendons == 0: - return wp.zeros((self._num_instances, 0), dtype=wp.float32, device=self.device) + return ProxyArray(wp.zeros((self._num_instances, 0), dtype=wp.float32, device=self.device)) if self._fixed_tendon_offset is None: - return wp.zeros((self._num_instances, self._num_fixed_tendons), dtype=wp.float32, device=self.device) - return self._fixed_tendon_offset + self._fixed_tendon_offset = wp.zeros( + (self._num_instances, self._num_fixed_tendons), dtype=wp.float32, device=self.device + ) + self._fixed_tendon_offset_ta = None + if self._fixed_tendon_offset_ta is None: + self._fixed_tendon_offset_ta = ProxyArray(self._fixed_tendon_offset) + return self._fixed_tendon_offset_ta @property - def fixed_tendon_pos_limits(self) -> wp.array: + def fixed_tendon_pos_limits(self) -> ProxyArray: """Fixed tendon position limits. dtype=wp.vec2f, shape: (N, num_fixed_tendons).""" if self._num_fixed_tendons == 0: - return wp.zeros((self._num_instances, 0, 2), dtype=wp.float32, device=self.device) + return ProxyArray(wp.zeros((self._num_instances, 0, 2), dtype=wp.float32, device=self.device)) if self._fixed_tendon_pos_limits is None: np_limits = np.zeros((self._num_instances, self._num_fixed_tendons, 2), dtype=np.float32) np_limits[..., 0] = -float("inf") np_limits[..., 1] = float("inf") - return wp.array(np_limits, dtype=wp.float32, device=self.device).view(wp.vec2f) - return self._fixed_tendon_pos_limits + self._fixed_tendon_pos_limits = wp.array(np_limits, dtype=wp.float32, device=self.device).view(wp.vec2f) + self._fixed_tendon_pos_limits_ta = None + if self._fixed_tendon_pos_limits_ta is None: + self._fixed_tendon_pos_limits_ta = ProxyArray(self._fixed_tendon_pos_limits) + return self._fixed_tendon_pos_limits_ta # -- Spatial tendon properties -- @property - def spatial_tendon_stiffness(self) -> wp.array: + def spatial_tendon_stiffness(self) -> ProxyArray: """Spatial tendon stiffness. Shape: (N, num_spatial_tendons).""" if self._num_spatial_tendons == 0: - return wp.zeros((self._num_instances, 0), dtype=wp.float32, device=self.device) + return ProxyArray(wp.zeros((self._num_instances, 0), dtype=wp.float32, device=self.device)) if self._spatial_tendon_stiffness is None: - return wp.zeros((self._num_instances, self._num_spatial_tendons), dtype=wp.float32, device=self.device) - return self._spatial_tendon_stiffness + self._spatial_tendon_stiffness = wp.zeros( + (self._num_instances, self._num_spatial_tendons), dtype=wp.float32, device=self.device + ) + self._spatial_tendon_stiffness_ta = None + if self._spatial_tendon_stiffness_ta is None: + self._spatial_tendon_stiffness_ta = ProxyArray(self._spatial_tendon_stiffness) + return self._spatial_tendon_stiffness_ta @property - def spatial_tendon_damping(self) -> wp.array: + def spatial_tendon_damping(self) -> ProxyArray: """Spatial tendon damping. Shape: (N, num_spatial_tendons).""" if self._num_spatial_tendons == 0: - return wp.zeros((self._num_instances, 0), dtype=wp.float32, device=self.device) + return ProxyArray(wp.zeros((self._num_instances, 0), dtype=wp.float32, device=self.device)) if self._spatial_tendon_damping is None: - return wp.zeros((self._num_instances, self._num_spatial_tendons), dtype=wp.float32, device=self.device) - return self._spatial_tendon_damping + self._spatial_tendon_damping = wp.zeros( + (self._num_instances, self._num_spatial_tendons), dtype=wp.float32, device=self.device + ) + self._spatial_tendon_damping_ta = None + if self._spatial_tendon_damping_ta is None: + self._spatial_tendon_damping_ta = ProxyArray(self._spatial_tendon_damping) + return self._spatial_tendon_damping_ta @property - def spatial_tendon_limit_stiffness(self) -> wp.array: + def spatial_tendon_limit_stiffness(self) -> ProxyArray: """Spatial tendon limit stiffness. Shape: (N, num_spatial_tendons).""" if self._num_spatial_tendons == 0: - return wp.zeros((self._num_instances, 0), dtype=wp.float32, device=self.device) + return ProxyArray(wp.zeros((self._num_instances, 0), dtype=wp.float32, device=self.device)) if self._spatial_tendon_limit_stiffness is None: - return wp.zeros((self._num_instances, self._num_spatial_tendons), dtype=wp.float32, device=self.device) - return self._spatial_tendon_limit_stiffness + self._spatial_tendon_limit_stiffness = wp.zeros( + (self._num_instances, self._num_spatial_tendons), dtype=wp.float32, device=self.device + ) + self._spatial_tendon_limit_stiffness_ta = None + if self._spatial_tendon_limit_stiffness_ta is None: + self._spatial_tendon_limit_stiffness_ta = ProxyArray(self._spatial_tendon_limit_stiffness) + return self._spatial_tendon_limit_stiffness_ta @property - def spatial_tendon_offset(self) -> wp.array: + def spatial_tendon_offset(self) -> ProxyArray: """Spatial tendon offset. Shape: (N, num_spatial_tendons).""" if self._num_spatial_tendons == 0: - return wp.zeros((self._num_instances, 0), dtype=wp.float32, device=self.device) + return ProxyArray(wp.zeros((self._num_instances, 0), dtype=wp.float32, device=self.device)) if self._spatial_tendon_offset is None: - return wp.zeros((self._num_instances, self._num_spatial_tendons), dtype=wp.float32, device=self.device) - return self._spatial_tendon_offset + self._spatial_tendon_offset = wp.zeros( + (self._num_instances, self._num_spatial_tendons), dtype=wp.float32, device=self.device + ) + self._spatial_tendon_offset_ta = None + if self._spatial_tendon_offset_ta is None: + self._spatial_tendon_offset_ta = ProxyArray(self._spatial_tendon_offset) + return self._spatial_tendon_offset_ta # -- Update method -- @@ -786,138 +1042,183 @@ def _identity_quat_np(self, *shape: int) -> np.ndarray: def set_default_root_pose(self, value: torch.Tensor) -> None: self._default_root_pose = wp.from_torch(value.to(self.device).contiguous()).view(wp.transformf) + self._default_root_pose_ta = None def set_default_root_vel(self, value: torch.Tensor) -> None: self._default_root_vel = wp.from_torch(value.to(self.device).contiguous()) + self._default_root_vel_ta = None def set_default_joint_pos(self, value: torch.Tensor) -> None: self._default_joint_pos = wp.from_torch(value.to(self.device).contiguous()) + self._default_joint_pos_ta = None def set_default_joint_vel(self, value: torch.Tensor) -> None: self._default_joint_vel = wp.from_torch(value.to(self.device).contiguous()) + self._default_joint_vel_ta = None def set_joint_pos_target(self, value: torch.Tensor) -> None: self._joint_pos_target = wp.from_torch(value.to(self.device).contiguous()) + self._joint_pos_target_ta = None def set_joint_vel_target(self, value: torch.Tensor) -> None: self._joint_vel_target = wp.from_torch(value.to(self.device).contiguous()) + self._joint_vel_target_ta = None def set_joint_effort_target(self, value: torch.Tensor) -> None: self._joint_effort_target = wp.from_torch(value.to(self.device).contiguous()) + self._joint_effort_target_ta = None def set_computed_torque(self, value: torch.Tensor) -> None: self._computed_torque = wp.from_torch(value.to(self.device).contiguous()) + self._computed_torque_ta = None def set_applied_torque(self, value: torch.Tensor) -> None: self._applied_torque = wp.from_torch(value.to(self.device).contiguous()) + self._applied_torque_ta = None def set_joint_stiffness(self, value: torch.Tensor) -> None: self._joint_stiffness = wp.from_torch(value.to(self.device).contiguous()) + self._joint_stiffness_ta = None def set_joint_damping(self, value: torch.Tensor) -> None: self._joint_damping = wp.from_torch(value.to(self.device).contiguous()) + self._joint_damping_ta = None def set_joint_armature(self, value: torch.Tensor) -> None: self._joint_armature = wp.from_torch(value.to(self.device).contiguous()) + self._joint_armature_ta = None def set_joint_friction_coeff(self, value: torch.Tensor) -> None: self._joint_friction_coeff = wp.from_torch(value.to(self.device).contiguous()) + self._joint_friction_coeff_ta = None def set_joint_pos_limits(self, value: torch.Tensor) -> None: self._joint_pos_limits = wp.from_torch(value.to(self.device).contiguous()) + self._joint_pos_limits_ta = None def set_joint_vel_limits(self, value: torch.Tensor) -> None: self._joint_vel_limits = wp.from_torch(value.to(self.device).contiguous()) + self._joint_vel_limits_ta = None def set_joint_effort_limits(self, value: torch.Tensor) -> None: self._joint_effort_limits = wp.from_torch(value.to(self.device).contiguous()) + self._joint_effort_limits_ta = None def set_soft_joint_pos_limits(self, value: torch.Tensor) -> None: self._soft_joint_pos_limits = wp.from_torch(value.to(self.device).contiguous()) + self._soft_joint_pos_limits_ta = None def set_soft_joint_vel_limits(self, value: torch.Tensor) -> None: self._soft_joint_vel_limits = wp.from_torch(value.to(self.device).contiguous()) + self._soft_joint_vel_limits_ta = None def set_gear_ratio(self, value: torch.Tensor) -> None: self._gear_ratio = wp.from_torch(value.to(self.device).contiguous()) + self._gear_ratio_ta = None def set_joint_pos(self, value: torch.Tensor) -> None: self._joint_pos = wp.from_torch(value.to(self.device).contiguous()) + self._joint_pos_ta = None def set_joint_vel(self, value: torch.Tensor) -> None: self._joint_vel = wp.from_torch(value.to(self.device).contiguous()) + self._joint_vel_ta = None def set_joint_acc(self, value: torch.Tensor) -> None: self._joint_acc = wp.from_torch(value.to(self.device).contiguous()) + self._joint_acc_ta = None def set_root_link_pose_w(self, value: torch.Tensor) -> None: self._root_link_pose_w = wp.from_torch(value.to(self.device).contiguous()).view(wp.transformf) + self._root_link_pose_w_ta = None def set_root_link_vel_w(self, value: torch.Tensor) -> None: self._root_link_vel_w = wp.from_torch(value.to(self.device).contiguous()) + self._root_link_vel_w_ta = None def set_root_com_pose_w(self, value: torch.Tensor) -> None: self._root_com_pose_w = wp.from_torch(value.to(self.device).contiguous()).view(wp.transformf) + self._root_com_pose_w_ta = None def set_root_com_vel_w(self, value: torch.Tensor) -> None: self._root_com_vel_w = wp.from_torch(value.to(self.device).contiguous()) + self._root_com_vel_w_ta = None def set_body_link_pose_w(self, value: torch.Tensor) -> None: self._body_link_pose_w = wp.from_torch(value.to(self.device).contiguous()).view(wp.transformf) + self._body_link_pose_w_ta = None def set_body_link_vel_w(self, value: torch.Tensor) -> None: self._body_link_vel_w = wp.from_torch(value.to(self.device).contiguous()) + self._body_link_vel_w_ta = None def set_body_com_pose_w(self, value: torch.Tensor) -> None: self._body_com_pose_w = wp.from_torch(value.to(self.device).contiguous()).view(wp.transformf) + self._body_com_pose_w_ta = None def set_body_com_vel_w(self, value: torch.Tensor) -> None: self._body_com_vel_w = wp.from_torch(value.to(self.device).contiguous()) + self._body_com_vel_w_ta = None def set_body_com_acc_w(self, value: torch.Tensor) -> None: self._body_com_acc_w = wp.from_torch(value.to(self.device).contiguous()) + self._body_com_acc_w_ta = None def set_body_com_pose_b(self, value: torch.Tensor) -> None: self._body_com_pose_b = wp.from_torch(value.to(self.device).contiguous()).view(wp.transformf) + self._body_com_pose_b_ta = None def set_body_mass(self, value: torch.Tensor) -> None: self._body_mass = wp.from_torch(value.to(self.device).contiguous()) + self._body_mass_ta = None def set_body_inertia(self, value: torch.Tensor) -> None: self._body_inertia = wp.from_torch(value.to(self.device).contiguous()) + self._body_inertia_ta = None def set_body_incoming_joint_wrench_b(self, value: torch.Tensor) -> None: self._body_incoming_joint_wrench_b = wp.from_torch(value.to(self.device).contiguous()) + self._body_incoming_joint_wrench_b_ta = None def set_fixed_tendon_stiffness(self, value: torch.Tensor) -> None: self._fixed_tendon_stiffness = wp.from_torch(value.to(self.device).contiguous()) + self._fixed_tendon_stiffness_ta = None def set_fixed_tendon_damping(self, value: torch.Tensor) -> None: self._fixed_tendon_damping = wp.from_torch(value.to(self.device).contiguous()) + self._fixed_tendon_damping_ta = None def set_fixed_tendon_limit_stiffness(self, value: torch.Tensor) -> None: self._fixed_tendon_limit_stiffness = wp.from_torch(value.to(self.device).contiguous()) + self._fixed_tendon_limit_stiffness_ta = None def set_fixed_tendon_rest_length(self, value: torch.Tensor) -> None: self._fixed_tendon_rest_length = wp.from_torch(value.to(self.device).contiguous()) + self._fixed_tendon_rest_length_ta = None def set_fixed_tendon_offset(self, value: torch.Tensor) -> None: self._fixed_tendon_offset = wp.from_torch(value.to(self.device).contiguous()) + self._fixed_tendon_offset_ta = None def set_fixed_tendon_pos_limits(self, value: torch.Tensor) -> None: self._fixed_tendon_pos_limits = wp.from_torch(value.to(self.device).contiguous()) + self._fixed_tendon_pos_limits_ta = None def set_spatial_tendon_stiffness(self, value: torch.Tensor) -> None: self._spatial_tendon_stiffness = wp.from_torch(value.to(self.device).contiguous()) + self._spatial_tendon_stiffness_ta = None def set_spatial_tendon_damping(self, value: torch.Tensor) -> None: self._spatial_tendon_damping = wp.from_torch(value.to(self.device).contiguous()) + self._spatial_tendon_damping_ta = None def set_spatial_tendon_limit_stiffness(self, value: torch.Tensor) -> None: self._spatial_tendon_limit_stiffness = wp.from_torch(value.to(self.device).contiguous()) + self._spatial_tendon_limit_stiffness_ta = None def set_spatial_tendon_offset(self, value: torch.Tensor) -> None: self._spatial_tendon_offset = wp.from_torch(value.to(self.device).contiguous()) + self._spatial_tendon_offset_ta = None def set_mock_data(self, **kwargs) -> None: """Bulk setter for mock data. @@ -1415,6 +1716,7 @@ def set_joint_position_target( pos_target = wp.to_torch(self._data._joint_pos_target) pos_target[env_ids, joint_ids] = target.to(self._device) self._data._joint_pos_target = wp.from_torch(pos_target) + self._data._joint_pos_target_ta = None def set_joint_velocity_target( self, @@ -1434,6 +1736,7 @@ def set_joint_velocity_target( vel_target = wp.to_torch(self._data._joint_vel_target) vel_target[env_ids, joint_ids] = target.to(self._device) self._data._joint_vel_target = wp.from_torch(vel_target) + self._data._joint_vel_target_ta = None def set_joint_effort_target( self, @@ -1453,6 +1756,7 @@ def set_joint_effort_target( effort_target = wp.to_torch(self._data._joint_effort_target) effort_target[env_ids, joint_ids] = target.to(self._device) self._data._joint_effort_target = wp.from_torch(effort_target) + self._data._joint_effort_target_ta = None # -- Tendon methods (fixed) -- 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 index 9d2543e14bdc..32f5f2bc3017 100644 --- a/source/isaaclab/isaaclab/test/mock_interfaces/assets/mock_rigid_object.py +++ b/source/isaaclab/isaaclab/test/mock_interfaces/assets/mock_rigid_object.py @@ -14,6 +14,8 @@ import torch import warp as wp +from isaaclab.utils.warp import ProxyArray + try: from isaaclab.assets.rigid_object.base_rigid_object_data import BaseRigidObjectData except (ImportError, ModuleNotFoundError): @@ -83,6 +85,22 @@ def __init__( self._body_mass: wp.array | None = None self._body_inertia: wp.array | None = None + # ProxyArray caches + self._default_root_pose_ta: ProxyArray | None = None + self._default_root_vel_ta: ProxyArray | None = None + self._root_link_pose_w_ta: ProxyArray | None = None + self._root_link_vel_w_ta: ProxyArray | None = None + self._root_com_pose_w_ta: ProxyArray | None = None + self._root_com_vel_w_ta: ProxyArray | None = None + self._body_link_pose_w_ta: ProxyArray | None = None + self._body_link_vel_w_ta: ProxyArray | None = None + self._body_com_pose_w_ta: ProxyArray | None = None + self._body_com_vel_w_ta: ProxyArray | None = None + self._body_com_acc_w_ta: ProxyArray | None = None + self._body_com_pose_b_ta: ProxyArray | None = None + self._body_mass_ta: ProxyArray | None = None + self._body_inertia_ta: ProxyArray | None = None + def _identity_quat(self, *shape: int) -> wp.array: """Create identity quaternion warp array (w, x, y, z) = (1, 0, 0, 0).""" quat_np = np.zeros((*shape, 4), dtype=np.float32) @@ -92,342 +110,404 @@ def _identity_quat(self, *shape: int) -> wp.array: # -- Default state properties -- @property - def default_root_pose(self) -> wp.array: + def default_root_pose(self) -> ProxyArray: """Default root pose. dtype=wp.transformf, shape: (N,).""" if self._default_root_pose is None: pose_np = np.zeros((self._num_instances, 7), dtype=np.float32) pose_np[..., 6] = 1.0 # identity quat qw=1, transformf layout: (px,py,pz,qx,qy,qz,qw) self._default_root_pose = wp.array(pose_np, dtype=wp.float32, device=self.device).view(wp.transformf) - return self._default_root_pose + self._default_root_pose_ta = None + if self._default_root_pose_ta is None: + self._default_root_pose_ta = ProxyArray(self._default_root_pose) + return self._default_root_pose_ta @property - def default_root_vel(self) -> wp.array: + def default_root_vel(self) -> ProxyArray: """Default root velocity. dtype=wp.spatial_vectorf, shape: (N,).""" if self._default_root_vel is None: self._default_root_vel = wp.zeros((self._num_instances, 6), dtype=wp.float32, device=self.device).view( wp.spatial_vectorf ) - return self._default_root_vel + self._default_root_vel_ta = None + if self._default_root_vel_ta is None: + self._default_root_vel_ta = ProxyArray(self._default_root_vel) + return self._default_root_vel_ta @property - def default_root_state(self) -> wp.array: + def default_root_state(self) -> ProxyArray: """Default root state. Shape: (N, 13).""" - pose = wp.to_torch(self.default_root_pose) - vel = wp.to_torch(self.default_root_vel) - return wp.from_torch(torch.cat([pose, vel], dim=-1)) + pose = self.default_root_pose.torch + vel = self.default_root_vel.torch + return ProxyArray(wp.from_torch(torch.cat([pose, vel], dim=-1))) # -- Root state properties (link frame) -- @property - def root_link_pose_w(self) -> wp.array: + def root_link_pose_w(self) -> ProxyArray: """Root link pose in world frame. dtype=wp.transformf, shape: (N,).""" if self._root_link_pose_w is None: pose_np = np.zeros((self._num_instances, 7), dtype=np.float32) pose_np[..., 6] = 1.0 # identity quat qw=1, transformf layout: (px,py,pz,qx,qy,qz,qw) self._root_link_pose_w = wp.array(pose_np, dtype=wp.float32, device=self.device).view(wp.transformf) - return self._root_link_pose_w + self._root_link_pose_w_ta = None + if self._root_link_pose_w_ta is None: + self._root_link_pose_w_ta = ProxyArray(self._root_link_pose_w) + return self._root_link_pose_w_ta @property - def root_link_vel_w(self) -> wp.array: + def root_link_vel_w(self) -> ProxyArray: """Root link velocity in world frame. dtype=wp.spatial_vectorf, shape: (N,).""" if self._root_link_vel_w is None: self._root_link_vel_w = wp.zeros((self._num_instances, 6), dtype=wp.float32, device=self.device).view( wp.spatial_vectorf ) - return self._root_link_vel_w + self._root_link_vel_w_ta = None + if self._root_link_vel_w_ta is None: + self._root_link_vel_w_ta = ProxyArray(self._root_link_vel_w) + return self._root_link_vel_w_ta @property - def root_link_state_w(self) -> wp.array: + def root_link_state_w(self) -> ProxyArray: """Root link state in world frame. Shape: (N, 13).""" - pose = wp.to_torch(self.root_link_pose_w) - vel = wp.to_torch(self.root_link_vel_w) - return wp.from_torch(torch.cat([pose, vel], dim=-1)) + pose = self.root_link_pose_w.torch + vel = self.root_link_vel_w.torch + return ProxyArray(wp.from_torch(torch.cat([pose, vel], dim=-1))) # Sliced properties (zero-copy pointer arithmetic on transformf) @property - def root_link_pos_w(self) -> wp.array: + def root_link_pos_w(self) -> ProxyArray: """Root link position. Shape: (N,), dtype=wp.vec3f.""" - t = self.root_link_pose_w - return wp.array(ptr=t.ptr, shape=t.shape, dtype=wp.vec3f, strides=t.strides, device=self.device) + t = self.root_link_pose_w.warp + return ProxyArray(wp.array(ptr=t.ptr, shape=t.shape, dtype=wp.vec3f, strides=t.strides, device=self.device)) # Sliced properties (convert through torch for simplicity in mock) @property - def root_link_quat_w(self) -> wp.array: + def root_link_quat_w(self) -> ProxyArray: """Root link orientation. Shape: (N,), dtype=wp.quatf.""" - t = self.root_link_pose_w - return wp.array(ptr=t.ptr + 3 * 4, shape=t.shape, dtype=wp.quatf, strides=t.strides, device=self.device) + t = self.root_link_pose_w.warp + return ProxyArray( + wp.array(ptr=t.ptr + 3 * 4, shape=t.shape, dtype=wp.quatf, strides=t.strides, device=self.device) + ) @property - def root_link_lin_vel_w(self) -> wp.array: + def root_link_lin_vel_w(self) -> ProxyArray: """Root link linear velocity. Shape: (N,), dtype=wp.vec3f.""" - v = self.root_link_vel_w - return wp.array(ptr=v.ptr, shape=v.shape, dtype=wp.vec3f, strides=v.strides, device=self.device) + v = self.root_link_vel_w.warp + return ProxyArray(wp.array(ptr=v.ptr, shape=v.shape, dtype=wp.vec3f, strides=v.strides, device=self.device)) @property - def root_link_ang_vel_w(self) -> wp.array: + def root_link_ang_vel_w(self) -> ProxyArray: """Root link angular velocity. Shape: (N,), dtype=wp.vec3f.""" - v = self.root_link_vel_w - return wp.array(ptr=v.ptr + 3 * 4, shape=v.shape, dtype=wp.vec3f, strides=v.strides, device=self.device) + v = self.root_link_vel_w.warp + return ProxyArray( + wp.array(ptr=v.ptr + 3 * 4, shape=v.shape, dtype=wp.vec3f, strides=v.strides, device=self.device) + ) # -- Root state properties (CoM frame) -- @property - def root_com_pose_w(self) -> wp.array: + def root_com_pose_w(self) -> ProxyArray: """Root CoM pose in world frame. dtype=wp.transformf, shape: (N,).""" if self._root_com_pose_w is None: - self._root_com_pose_w = wp.clone(self.root_link_pose_w, self.device) - return self._root_com_pose_w + self._root_com_pose_w = wp.clone(self.root_link_pose_w.warp, self.device) + self._root_com_pose_w_ta = None + if self._root_com_pose_w_ta is None: + self._root_com_pose_w_ta = ProxyArray(self._root_com_pose_w) + return self._root_com_pose_w_ta @property - def root_com_vel_w(self) -> wp.array: + def root_com_vel_w(self) -> ProxyArray: """Root CoM velocity in world frame. dtype=wp.spatial_vectorf, shape: (N,).""" if self._root_com_vel_w is None: - self._root_com_vel_w = wp.clone(self.root_link_vel_w, self.device) - return self._root_com_vel_w + self._root_com_vel_w = wp.clone(self.root_link_vel_w.warp, self.device) + self._root_com_vel_w_ta = None + if self._root_com_vel_w_ta is None: + self._root_com_vel_w_ta = ProxyArray(self._root_com_vel_w) + return self._root_com_vel_w_ta @property - def root_com_state_w(self) -> wp.array: + def root_com_state_w(self) -> ProxyArray: """Root CoM state in world frame. Shape: (N, 13).""" - pose = wp.to_torch(self.root_com_pose_w) - vel = wp.to_torch(self.root_com_vel_w) - return wp.from_torch(torch.cat([pose, vel], dim=-1)) + pose = self.root_com_pose_w.torch + vel = self.root_com_vel_w.torch + return ProxyArray(wp.from_torch(torch.cat([pose, vel], dim=-1))) @property - def root_state_w(self) -> wp.array: + def root_state_w(self) -> ProxyArray: """Root state (link pose + CoM velocity). Shape: (N, 13).""" - pose = wp.to_torch(self.root_link_pose_w) - vel = wp.to_torch(self.root_com_vel_w) - return wp.from_torch(torch.cat([pose, vel], dim=-1)) + pose = self.root_link_pose_w.torch + vel = self.root_com_vel_w.torch + return ProxyArray(wp.from_torch(torch.cat([pose, vel], dim=-1))) # Sliced properties (zero-copy pointer arithmetic on transformf) @property - def root_com_pos_w(self) -> wp.array: + def root_com_pos_w(self) -> ProxyArray: """Root CoM position. Shape: (N,), dtype=wp.vec3f.""" - t = self.root_com_pose_w - return wp.array(ptr=t.ptr, shape=t.shape, dtype=wp.vec3f, strides=t.strides, device=self.device) + t = self.root_com_pose_w.warp + return ProxyArray(wp.array(ptr=t.ptr, shape=t.shape, dtype=wp.vec3f, strides=t.strides, device=self.device)) # Sliced properties (convert through torch for simplicity in mock) @property - def root_com_quat_w(self) -> wp.array: + def root_com_quat_w(self) -> ProxyArray: """Root CoM orientation. Shape: (N,), dtype=wp.quatf.""" - t = self.root_com_pose_w - return wp.array(ptr=t.ptr + 3 * 4, shape=t.shape, dtype=wp.quatf, strides=t.strides, device=self.device) + t = self.root_com_pose_w.warp + return ProxyArray( + wp.array(ptr=t.ptr + 3 * 4, shape=t.shape, dtype=wp.quatf, strides=t.strides, device=self.device) + ) @property - def root_com_lin_vel_w(self) -> wp.array: + def root_com_lin_vel_w(self) -> ProxyArray: """Root CoM linear velocity. Shape: (N,), dtype=wp.vec3f.""" - v = self.root_com_vel_w - return wp.array(ptr=v.ptr, shape=v.shape, dtype=wp.vec3f, strides=v.strides, device=self.device) + v = self.root_com_vel_w.warp + return ProxyArray(wp.array(ptr=v.ptr, shape=v.shape, dtype=wp.vec3f, strides=v.strides, device=self.device)) @property - def root_com_ang_vel_w(self) -> wp.array: + def root_com_ang_vel_w(self) -> ProxyArray: """Root CoM angular velocity. Shape: (N,), dtype=wp.vec3f.""" - v = self.root_com_vel_w - return wp.array(ptr=v.ptr + 3 * 4, shape=v.shape, dtype=wp.vec3f, strides=v.strides, device=self.device) + v = self.root_com_vel_w.warp + return ProxyArray( + wp.array(ptr=v.ptr + 3 * 4, shape=v.shape, dtype=wp.vec3f, strides=v.strides, device=self.device) + ) # -- Body state properties (link frame) -- @property - def body_link_pose_w(self) -> wp.array: + def body_link_pose_w(self) -> ProxyArray: """Body link pose in world frame. dtype=wp.transformf, shape: (N, 1).""" if self._body_link_pose_w is None: pose_np = np.zeros((self._num_instances, 1, 7), dtype=np.float32) pose_np[..., 6] = 1.0 # identity quat qw=1, transformf layout: (px,py,pz,qx,qy,qz,qw) self._body_link_pose_w = wp.array(pose_np, dtype=wp.float32, device=self.device).view(wp.transformf) - return self._body_link_pose_w + self._body_link_pose_w_ta = None + if self._body_link_pose_w_ta is None: + self._body_link_pose_w_ta = ProxyArray(self._body_link_pose_w) + return self._body_link_pose_w_ta @property - def body_link_vel_w(self) -> wp.array: + def body_link_vel_w(self) -> ProxyArray: """Body link velocity in world frame. dtype=wp.spatial_vectorf, shape: (N, 1).""" if self._body_link_vel_w is None: self._body_link_vel_w = wp.zeros((self._num_instances, 1, 6), dtype=wp.float32, device=self.device).view( wp.spatial_vectorf ) - return self._body_link_vel_w + self._body_link_vel_w_ta = None + if self._body_link_vel_w_ta is None: + self._body_link_vel_w_ta = ProxyArray(self._body_link_vel_w) + return self._body_link_vel_w_ta @property - def body_link_state_w(self) -> wp.array: + def body_link_state_w(self) -> ProxyArray: """Body link state in world frame. Shape: (N, 1, 13).""" - pose = wp.to_torch(self.body_link_pose_w) - vel = wp.to_torch(self.body_link_vel_w) - return wp.from_torch(torch.cat([pose, vel], dim=-1)) + pose = self.body_link_pose_w.torch + vel = self.body_link_vel_w.torch + return ProxyArray(wp.from_torch(torch.cat([pose, vel], dim=-1))) # Sliced properties (zero-copy pointer arithmetic on transformf) @property - def body_link_pos_w(self) -> wp.array: + def body_link_pos_w(self) -> ProxyArray: """Body link position. Shape: (N, 1), dtype=wp.vec3f.""" - t = self.body_link_pose_w - return wp.array(ptr=t.ptr, shape=t.shape, dtype=wp.vec3f, strides=t.strides, device=self.device) + t = self.body_link_pose_w.warp + return ProxyArray(wp.array(ptr=t.ptr, shape=t.shape, dtype=wp.vec3f, strides=t.strides, device=self.device)) # Sliced properties (convert through torch for simplicity in mock) @property - def body_link_quat_w(self) -> wp.array: + def body_link_quat_w(self) -> ProxyArray: """Body link orientation. Shape: (N, 1), dtype=wp.quatf.""" - t = self.body_link_pose_w - return wp.array(ptr=t.ptr + 3 * 4, shape=t.shape, dtype=wp.quatf, strides=t.strides, device=self.device) + t = self.body_link_pose_w.warp + return ProxyArray( + wp.array(ptr=t.ptr + 3 * 4, shape=t.shape, dtype=wp.quatf, strides=t.strides, device=self.device) + ) @property - def body_link_lin_vel_w(self) -> wp.array: + def body_link_lin_vel_w(self) -> ProxyArray: """Body link linear velocity. Shape: (N, 1), dtype=wp.vec3f.""" - v = self.body_link_vel_w - return wp.array(ptr=v.ptr, shape=v.shape, dtype=wp.vec3f, strides=v.strides, device=self.device) + v = self.body_link_vel_w.warp + return ProxyArray(wp.array(ptr=v.ptr, shape=v.shape, dtype=wp.vec3f, strides=v.strides, device=self.device)) @property - def body_link_ang_vel_w(self) -> wp.array: + def body_link_ang_vel_w(self) -> ProxyArray: """Body link angular velocity. Shape: (N, 1), dtype=wp.vec3f.""" - v = self.body_link_vel_w - return wp.array(ptr=v.ptr + 3 * 4, shape=v.shape, dtype=wp.vec3f, strides=v.strides, device=self.device) + v = self.body_link_vel_w.warp + return ProxyArray( + wp.array(ptr=v.ptr + 3 * 4, shape=v.shape, dtype=wp.vec3f, strides=v.strides, device=self.device) + ) # -- Body state properties (CoM frame) -- @property - def body_com_pose_w(self) -> wp.array: + def body_com_pose_w(self) -> ProxyArray: """Body CoM pose in world frame. dtype=wp.transformf, shape: (N, 1).""" if self._body_com_pose_w is None: - self._body_com_pose_w = wp.clone(self.body_link_pose_w, self.device) - return self._body_com_pose_w + self._body_com_pose_w = wp.clone(self.body_link_pose_w.warp, self.device) + self._body_com_pose_w_ta = None + if self._body_com_pose_w_ta is None: + self._body_com_pose_w_ta = ProxyArray(self._body_com_pose_w) + return self._body_com_pose_w_ta @property - def body_com_vel_w(self) -> wp.array: + def body_com_vel_w(self) -> ProxyArray: """Body CoM velocity in world frame. dtype=wp.spatial_vectorf, shape: (N, 1).""" if self._body_com_vel_w is None: - self._body_com_vel_w = wp.clone(self.body_link_vel_w, self.device) - return self._body_com_vel_w + self._body_com_vel_w = wp.clone(self.body_link_vel_w.warp, self.device) + self._body_com_vel_w_ta = None + if self._body_com_vel_w_ta is None: + self._body_com_vel_w_ta = ProxyArray(self._body_com_vel_w) + return self._body_com_vel_w_ta @property - def body_com_state_w(self) -> wp.array: + def body_com_state_w(self) -> ProxyArray: """Body CoM state in world frame. Shape: (N, 1, 13).""" - pose = wp.to_torch(self.body_com_pose_w) - vel = wp.to_torch(self.body_com_vel_w) - return wp.from_torch(torch.cat([pose, vel], dim=-1)) + pose = self.body_com_pose_w.torch + vel = self.body_com_vel_w.torch + return ProxyArray(wp.from_torch(torch.cat([pose, vel], dim=-1))) @property - def body_state_w(self) -> wp.array: + def body_state_w(self) -> ProxyArray: """Body state (link pose + CoM vel). Shape: (N, 1, 13).""" - pose = wp.to_torch(self.body_link_pose_w) - vel = wp.to_torch(self.body_com_vel_w) - return wp.from_torch(torch.cat([pose, vel], dim=-1)) + pose = self.body_link_pose_w.torch + vel = self.body_com_vel_w.torch + return ProxyArray(wp.from_torch(torch.cat([pose, vel], dim=-1))) @property - def body_com_acc_w(self) -> wp.array: + def body_com_acc_w(self) -> ProxyArray: """Body CoM acceleration in world frame. dtype=wp.spatial_vectorf, shape: (N, 1).""" if self._body_com_acc_w is None: self._body_com_acc_w = wp.zeros((self._num_instances, 1, 6), dtype=wp.float32, device=self.device).view( wp.spatial_vectorf ) - return self._body_com_acc_w + self._body_com_acc_w_ta = None + if self._body_com_acc_w_ta is None: + self._body_com_acc_w_ta = ProxyArray(self._body_com_acc_w) + return self._body_com_acc_w_ta @property - def body_com_pose_b(self) -> wp.array: + def body_com_pose_b(self) -> ProxyArray: """Body CoM pose in body frame. dtype=wp.transformf, shape: (N, 1).""" if self._body_com_pose_b is None: pose_np = np.zeros((self._num_instances, 1, 7), dtype=np.float32) pose_np[..., 6] = 1.0 # identity quat qw=1, transformf layout: (px,py,pz,qx,qy,qz,qw) self._body_com_pose_b = wp.array(pose_np, dtype=wp.float32, device=self.device).view(wp.transformf) - return self._body_com_pose_b + self._body_com_pose_b_ta = None + if self._body_com_pose_b_ta is None: + self._body_com_pose_b_ta = ProxyArray(self._body_com_pose_b) + return self._body_com_pose_b_ta # Sliced properties (zero-copy pointer arithmetic on transformf) @property - def body_com_pos_w(self) -> wp.array: + def body_com_pos_w(self) -> ProxyArray: """Body CoM position. Shape: (N, 1), dtype=wp.vec3f.""" - t = self.body_com_pose_w - return wp.array(ptr=t.ptr, shape=t.shape, dtype=wp.vec3f, strides=t.strides, device=self.device) + t = self.body_com_pose_w.warp + return ProxyArray(wp.array(ptr=t.ptr, shape=t.shape, dtype=wp.vec3f, strides=t.strides, device=self.device)) @property - def body_com_quat_w(self) -> wp.array: + def body_com_quat_w(self) -> ProxyArray: """Body CoM orientation. Shape: (N, 1), dtype=wp.quatf.""" - t = self.body_com_pose_w - return wp.array(ptr=t.ptr + 3 * 4, shape=t.shape, dtype=wp.quatf, strides=t.strides, device=self.device) + t = self.body_com_pose_w.warp + return ProxyArray( + wp.array(ptr=t.ptr + 3 * 4, shape=t.shape, dtype=wp.quatf, strides=t.strides, device=self.device) + ) @property - def body_com_lin_vel_w(self) -> wp.array: + def body_com_lin_vel_w(self) -> ProxyArray: """Body CoM linear velocity. Shape: (N, 1), dtype=wp.vec3f.""" - v = self.body_com_vel_w - return wp.array(ptr=v.ptr, shape=v.shape, dtype=wp.vec3f, strides=v.strides, device=self.device) + v = self.body_com_vel_w.warp + return ProxyArray(wp.array(ptr=v.ptr, shape=v.shape, dtype=wp.vec3f, strides=v.strides, device=self.device)) @property - def body_com_ang_vel_w(self) -> wp.array: + def body_com_ang_vel_w(self) -> ProxyArray: """Body CoM angular velocity. Shape: (N, 1), dtype=wp.vec3f.""" - v = self.body_com_vel_w - return wp.array(ptr=v.ptr + 3 * 4, shape=v.shape, dtype=wp.vec3f, strides=v.strides, device=self.device) + v = self.body_com_vel_w.warp + return ProxyArray( + wp.array(ptr=v.ptr + 3 * 4, shape=v.shape, dtype=wp.vec3f, strides=v.strides, device=self.device) + ) @property - def body_com_lin_acc_w(self) -> wp.array: + def body_com_lin_acc_w(self) -> ProxyArray: """Body CoM linear acceleration. Shape: (N, 1), dtype=wp.vec3f.""" - a = self.body_com_acc_w - return wp.array(ptr=a.ptr, shape=a.shape, dtype=wp.vec3f, strides=a.strides, device=self.device) + a = self.body_com_acc_w.warp + return ProxyArray(wp.array(ptr=a.ptr, shape=a.shape, dtype=wp.vec3f, strides=a.strides, device=self.device)) @property - def body_com_ang_acc_w(self) -> wp.array: + def body_com_ang_acc_w(self) -> ProxyArray: """Body CoM angular acceleration. Shape: (N, 1), dtype=wp.vec3f.""" - a = self.body_com_acc_w - return wp.array(ptr=a.ptr + 3 * 4, shape=a.shape, dtype=wp.vec3f, strides=a.strides, device=self.device) + a = self.body_com_acc_w.warp + return ProxyArray( + wp.array(ptr=a.ptr + 3 * 4, shape=a.shape, dtype=wp.vec3f, strides=a.strides, device=self.device) + ) @property - def body_com_pos_b(self) -> wp.array: + def body_com_pos_b(self) -> ProxyArray: """Body CoM position in body frame. Shape: (N, 1), dtype=wp.vec3f.""" - t = self.body_com_pose_b - return wp.array(ptr=t.ptr, shape=t.shape, dtype=wp.vec3f, strides=t.strides, device=self.device) + t = self.body_com_pose_b.warp + return ProxyArray(wp.array(ptr=t.ptr, shape=t.shape, dtype=wp.vec3f, strides=t.strides, device=self.device)) @property - def body_com_quat_b(self) -> wp.array: + def body_com_quat_b(self) -> ProxyArray: """Body CoM orientation in body frame. Shape: (N, 1), dtype=wp.quatf.""" - t = self.body_com_pose_b - return wp.array(ptr=t.ptr + 3 * 4, shape=t.shape, dtype=wp.quatf, strides=t.strides, device=self.device) + t = self.body_com_pose_b.warp + return ProxyArray( + wp.array(ptr=t.ptr + 3 * 4, shape=t.shape, dtype=wp.quatf, strides=t.strides, device=self.device) + ) # -- Body properties -- @property - def body_mass(self) -> wp.array: + def body_mass(self) -> ProxyArray: """Body mass. Shape: (N, 1).""" if self._body_mass is None: - return wp.ones((self._num_instances, 1), dtype=wp.float32, device=self.device) - return self._body_mass + self._body_mass = wp.ones((self._num_instances, 1), dtype=wp.float32, device=self.device) + self._body_mass_ta = None + if self._body_mass_ta is None: + self._body_mass_ta = ProxyArray(self._body_mass) + return self._body_mass_ta @property - def body_inertia(self) -> wp.array: + def body_inertia(self) -> ProxyArray: """Body inertia (flattened 3x3). Shape: (N, 1, 9).""" if self._body_inertia is None: inertia_np = np.zeros((self._num_instances, 1, 9), dtype=np.float32) inertia_np[..., 0] = 1.0 # Ixx inertia_np[..., 4] = 1.0 # Iyy inertia_np[..., 8] = 1.0 # Izz - return wp.array(inertia_np, dtype=wp.float32, device=self.device) - return self._body_inertia + self._body_inertia = wp.array(inertia_np, dtype=wp.float32, device=self.device) + self._body_inertia_ta = None + if self._body_inertia_ta is None: + self._body_inertia_ta = ProxyArray(self._body_inertia) + return self._body_inertia_ta # -- Derived properties -- @property - def projected_gravity_b(self) -> wp.array: + def projected_gravity_b(self) -> ProxyArray: """Gravity projection on body. Shape: (N,), dtype=wp.vec3f.""" gravity_np = np.zeros((self._num_instances, 3), dtype=np.float32) gravity_np[:, 2] = -1.0 - return wp.array(gravity_np, dtype=wp.float32, device=self.device).view(wp.vec3f) + return ProxyArray(wp.array(gravity_np, dtype=wp.float32, device=self.device).view(wp.vec3f)) @property - def heading_w(self) -> wp.array: + def heading_w(self) -> ProxyArray: """Yaw heading in world frame. Shape: (N,).""" - return wp.zeros((self._num_instances,), dtype=wp.float32, device=self.device) + return ProxyArray(wp.zeros((self._num_instances,), dtype=wp.float32, device=self.device)) @property - def root_link_lin_vel_b(self) -> wp.array: + def root_link_lin_vel_b(self) -> ProxyArray: """Root link linear velocity in body frame. Shape: (N,), dtype=wp.vec3f.""" - return wp.clone(self.root_link_lin_vel_w, self.device) + return ProxyArray(wp.clone(self.root_link_lin_vel_w.warp, self.device)) @property - def root_link_ang_vel_b(self) -> wp.array: + def root_link_ang_vel_b(self) -> ProxyArray: """Root link angular velocity in body frame. Shape: (N,), dtype=wp.vec3f.""" - return wp.clone(self.root_link_ang_vel_w, self.device) + return ProxyArray(wp.clone(self.root_link_ang_vel_w.warp, self.device)) @property - def root_com_lin_vel_b(self) -> wp.array: + def root_com_lin_vel_b(self) -> ProxyArray: """Root CoM linear velocity in body frame. Shape: (N,), dtype=wp.vec3f.""" - return wp.clone(self.root_com_lin_vel_w, self.device) + return ProxyArray(wp.clone(self.root_com_lin_vel_w.warp, self.device)) @property - def root_com_ang_vel_b(self) -> wp.array: + def root_com_ang_vel_b(self) -> ProxyArray: """Root CoM angular velocity in body frame. Shape: (N,), dtype=wp.vec3f.""" - return wp.clone(self.root_com_ang_vel_w, self.device) + return ProxyArray(wp.clone(self.root_com_ang_vel_w.warp, self.device)) # -- Shorthand properties (root_pose_w, body_pos_w, com_pos_b, etc.) -- @@ -449,45 +529,59 @@ def update(self, dt: float) -> None: def set_default_root_pose(self, value: torch.Tensor) -> None: self._default_root_pose = wp.from_torch(value.to(self.device).contiguous()).view(wp.transformf) + self._default_root_pose_ta = None def set_default_root_vel(self, value: torch.Tensor) -> None: self._default_root_vel = wp.from_torch(value.to(self.device).contiguous()) + self._default_root_vel_ta = None def set_root_link_pose_w(self, value: torch.Tensor) -> None: self._root_link_pose_w = wp.from_torch(value.to(self.device).contiguous()).view(wp.transformf) + self._root_link_pose_w_ta = None def set_root_link_vel_w(self, value: torch.Tensor) -> None: self._root_link_vel_w = wp.from_torch(value.to(self.device).contiguous()) + self._root_link_vel_w_ta = None def set_root_com_pose_w(self, value: torch.Tensor) -> None: self._root_com_pose_w = wp.from_torch(value.to(self.device).contiguous()).view(wp.transformf) + self._root_com_pose_w_ta = None def set_root_com_vel_w(self, value: torch.Tensor) -> None: self._root_com_vel_w = wp.from_torch(value.to(self.device).contiguous()) + self._root_com_vel_w_ta = None def set_body_link_pose_w(self, value: torch.Tensor) -> None: self._body_link_pose_w = wp.from_torch(value.to(self.device).contiguous()).view(wp.transformf) + self._body_link_pose_w_ta = None def set_body_link_vel_w(self, value: torch.Tensor) -> None: self._body_link_vel_w = wp.from_torch(value.to(self.device).contiguous()) + self._body_link_vel_w_ta = None def set_body_com_pose_w(self, value: torch.Tensor) -> None: self._body_com_pose_w = wp.from_torch(value.to(self.device).contiguous()).view(wp.transformf) + self._body_com_pose_w_ta = None def set_body_com_vel_w(self, value: torch.Tensor) -> None: self._body_com_vel_w = wp.from_torch(value.to(self.device).contiguous()) + self._body_com_vel_w_ta = None def set_body_com_acc_w(self, value: torch.Tensor) -> None: self._body_com_acc_w = wp.from_torch(value.to(self.device).contiguous()) + self._body_com_acc_w_ta = None def set_body_com_pose_b(self, value: torch.Tensor) -> None: self._body_com_pose_b = wp.from_torch(value.to(self.device).contiguous()).view(wp.transformf) + self._body_com_pose_b_ta = None def set_body_mass(self, value: torch.Tensor) -> None: self._body_mass = wp.from_torch(value.to(self.device).contiguous()) + self._body_mass_ta = None def set_body_inertia(self, value: torch.Tensor) -> None: self._body_inertia = wp.from_torch(value.to(self.device).contiguous()) + self._body_inertia_ta = None def set_mock_data(self, **kwargs) -> None: """Bulk setter for mock data.""" 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 index c9533ba2fc5c..0af998188a89 100644 --- 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 @@ -14,6 +14,8 @@ import torch import warp as wp +from isaaclab.utils.warp import ProxyArray + try: from isaaclab.assets.rigid_object_collection.base_rigid_object_collection_data import BaseRigidObjectCollectionData except (ImportError, ModuleNotFoundError): @@ -82,6 +84,18 @@ def __init__( self._body_mass: wp.array | None = None self._body_inertia: wp.array | None = None + # ProxyArray caches + self._default_body_pose_ta: ProxyArray | None = None + self._default_body_vel_ta: ProxyArray | None = None + self._body_link_pose_w_ta: ProxyArray | None = None + self._body_link_vel_w_ta: ProxyArray | None = None + self._body_com_pose_w_ta: ProxyArray | None = None + self._body_com_vel_w_ta: ProxyArray | None = None + self._body_com_acc_w_ta: ProxyArray | None = None + self._body_com_pose_b_ta: ProxyArray | None = None + self._body_mass_ta: ProxyArray | None = None + self._body_inertia_ta: ProxyArray | None = None + def _identity_quat(self, *shape: int) -> wp.array: """Create identity quaternion warp array (w, x, y, z) = (1, 0, 0, 0).""" quat_np = np.zeros((*shape, 4), dtype=np.float32) @@ -91,234 +105,276 @@ def _identity_quat(self, *shape: int) -> wp.array: # -- Default state properties -- @property - def default_body_pose(self) -> wp.array: + def default_body_pose(self) -> ProxyArray: """Default body poses. dtype=wp.transformf, shape: (N, num_bodies).""" if self._default_body_pose is None: pose_np = np.zeros((self._num_instances, self._num_bodies, 7), dtype=np.float32) pose_np[..., 6] = 1.0 # identity quat qw=1, transformf layout self._default_body_pose = wp.array(pose_np, dtype=wp.float32, device=self.device).view(wp.transformf) - return self._default_body_pose + self._default_body_pose_ta = None + if self._default_body_pose_ta is None: + self._default_body_pose_ta = ProxyArray(self._default_body_pose) + return self._default_body_pose_ta @property - def default_body_vel(self) -> wp.array: + def default_body_vel(self) -> ProxyArray: """Default body velocities. dtype=wp.spatial_vectorf, shape: (N, num_bodies).""" if self._default_body_vel is None: self._default_body_vel = wp.zeros( (self._num_instances, self._num_bodies, 6), dtype=wp.float32, device=self.device ).view(wp.spatial_vectorf) - return self._default_body_vel + self._default_body_vel_ta = None + if self._default_body_vel_ta is None: + self._default_body_vel_ta = ProxyArray(self._default_body_vel) + return self._default_body_vel_ta @property - def default_body_state(self) -> wp.array: + def default_body_state(self) -> ProxyArray: """Default body states. Shape: (N, num_bodies, 13).""" - pose = wp.to_torch(self.default_body_pose) - vel = wp.to_torch(self.default_body_vel) - return wp.from_torch(torch.cat([pose, vel], dim=-1)) + pose = self.default_body_pose.torch + vel = self.default_body_vel.torch + return ProxyArray(wp.from_torch(torch.cat([pose, vel], dim=-1))) # -- Body state properties (link frame) -- @property - def body_link_pose_w(self) -> wp.array: + def body_link_pose_w(self) -> ProxyArray: """Body link poses in world frame. dtype=wp.transformf, shape: (N, num_bodies).""" if self._body_link_pose_w is None: pose_np = np.zeros((self._num_instances, self._num_bodies, 7), dtype=np.float32) pose_np[..., 6] = 1.0 # identity quat qw=1, transformf layout: (px,py,pz,qx,qy,qz,qw) self._body_link_pose_w = wp.array(pose_np, dtype=wp.float32, device=self.device).view(wp.transformf) - return self._body_link_pose_w + self._body_link_pose_w_ta = None + if self._body_link_pose_w_ta is None: + self._body_link_pose_w_ta = ProxyArray(self._body_link_pose_w) + return self._body_link_pose_w_ta @property - def body_link_vel_w(self) -> wp.array: + def body_link_vel_w(self) -> ProxyArray: """Body link velocities in world frame. dtype=wp.spatial_vectorf, shape: (N, num_bodies).""" if self._body_link_vel_w is None: self._body_link_vel_w = wp.zeros( (self._num_instances, self._num_bodies, 6), dtype=wp.float32, device=self.device ).view(wp.spatial_vectorf) - return self._body_link_vel_w + self._body_link_vel_w_ta = None + if self._body_link_vel_w_ta is None: + self._body_link_vel_w_ta = ProxyArray(self._body_link_vel_w) + return self._body_link_vel_w_ta @property - def body_link_state_w(self) -> wp.array: + def body_link_state_w(self) -> ProxyArray: """Body link states in world frame. Shape: (N, num_bodies, 13).""" - pose = wp.to_torch(self.body_link_pose_w) - vel = wp.to_torch(self.body_link_vel_w) - return wp.from_torch(torch.cat([pose, vel], dim=-1)) + pose = self.body_link_pose_w.torch + vel = self.body_link_vel_w.torch + return ProxyArray(wp.from_torch(torch.cat([pose, vel], dim=-1))) # Sliced properties (zero-copy pointer arithmetic on transformf) @property - def body_link_pos_w(self) -> wp.array: + def body_link_pos_w(self) -> ProxyArray: """Body link positions. Shape: (N, num_bodies), dtype=wp.vec3f.""" - t = self.body_link_pose_w - return wp.array(ptr=t.ptr, shape=t.shape, dtype=wp.vec3f, strides=t.strides, device=self.device) + t = self.body_link_pose_w.warp + return ProxyArray(wp.array(ptr=t.ptr, shape=t.shape, dtype=wp.vec3f, strides=t.strides, device=self.device)) @property - def body_link_quat_w(self) -> wp.array: + def body_link_quat_w(self) -> ProxyArray: """Body link orientations. Shape: (N, num_bodies), dtype=wp.quatf.""" - t = self.body_link_pose_w - return wp.array(ptr=t.ptr + 3 * 4, shape=t.shape, dtype=wp.quatf, strides=t.strides, device=self.device) + t = self.body_link_pose_w.warp + return ProxyArray( + wp.array(ptr=t.ptr + 3 * 4, shape=t.shape, dtype=wp.quatf, strides=t.strides, device=self.device) + ) @property - def body_link_lin_vel_w(self) -> wp.array: + def body_link_lin_vel_w(self) -> ProxyArray: """Body link linear velocities. Shape: (N, num_bodies), dtype=wp.vec3f.""" - v = self.body_link_vel_w - return wp.array(ptr=v.ptr, shape=v.shape, dtype=wp.vec3f, strides=v.strides, device=self.device) + v = self.body_link_vel_w.warp + return ProxyArray(wp.array(ptr=v.ptr, shape=v.shape, dtype=wp.vec3f, strides=v.strides, device=self.device)) @property - def body_link_ang_vel_w(self) -> wp.array: + def body_link_ang_vel_w(self) -> ProxyArray: """Body link angular velocities. Shape: (N, num_bodies), dtype=wp.vec3f.""" - v = self.body_link_vel_w - return wp.array(ptr=v.ptr + 3 * 4, shape=v.shape, dtype=wp.vec3f, strides=v.strides, device=self.device) + v = self.body_link_vel_w.warp + return ProxyArray( + wp.array(ptr=v.ptr + 3 * 4, shape=v.shape, dtype=wp.vec3f, strides=v.strides, device=self.device) + ) # -- Body state properties (CoM frame) -- @property - def body_com_pose_w(self) -> wp.array: + def body_com_pose_w(self) -> ProxyArray: """Body CoM poses in world frame. dtype=wp.transformf, shape: (N, num_bodies).""" if self._body_com_pose_w is None: - self._body_com_pose_w = wp.clone(self.body_link_pose_w, self.device) - return self._body_com_pose_w + self._body_com_pose_w = wp.clone(self.body_link_pose_w.warp, self.device) + self._body_com_pose_w_ta = None + if self._body_com_pose_w_ta is None: + self._body_com_pose_w_ta = ProxyArray(self._body_com_pose_w) + return self._body_com_pose_w_ta @property - def body_com_vel_w(self) -> wp.array: + def body_com_vel_w(self) -> ProxyArray: """Body CoM velocities in world frame. dtype=wp.spatial_vectorf, shape: (N, num_bodies).""" if self._body_com_vel_w is None: - self._body_com_vel_w = wp.clone(self.body_link_vel_w, self.device) - return self._body_com_vel_w + self._body_com_vel_w = wp.clone(self.body_link_vel_w.warp, self.device) + self._body_com_vel_w_ta = None + if self._body_com_vel_w_ta is None: + self._body_com_vel_w_ta = ProxyArray(self._body_com_vel_w) + return self._body_com_vel_w_ta @property - def body_com_state_w(self) -> wp.array: + def body_com_state_w(self) -> ProxyArray: """Body CoM states in world frame. Shape: (N, num_bodies, 13).""" - pose = wp.to_torch(self.body_com_pose_w) - vel = wp.to_torch(self.body_com_vel_w) - return wp.from_torch(torch.cat([pose, vel], dim=-1)) + pose = self.body_com_pose_w.torch + vel = self.body_com_vel_w.torch + return ProxyArray(wp.from_torch(torch.cat([pose, vel], dim=-1))) @property - def body_com_acc_w(self) -> wp.array: + def body_com_acc_w(self) -> ProxyArray: """Body CoM accelerations in world frame. dtype=wp.spatial_vectorf, shape: (N, num_bodies).""" if self._body_com_acc_w is None: self._body_com_acc_w = wp.zeros( (self._num_instances, self._num_bodies, 6), dtype=wp.float32, device=self.device ).view(wp.spatial_vectorf) - return self._body_com_acc_w + self._body_com_acc_w_ta = None + if self._body_com_acc_w_ta is None: + self._body_com_acc_w_ta = ProxyArray(self._body_com_acc_w) + return self._body_com_acc_w_ta @property - def body_com_pose_b(self) -> wp.array: + def body_com_pose_b(self) -> ProxyArray: """Body CoM poses in body frame. dtype=wp.transformf, shape: (N, num_bodies).""" if self._body_com_pose_b is None: pose_np = np.zeros((self._num_instances, self._num_bodies, 7), dtype=np.float32) pose_np[..., 6] = 1.0 # identity quat qw=1, transformf layout: (px,py,pz,qx,qy,qz,qw) self._body_com_pose_b = wp.array(pose_np, dtype=wp.float32, device=self.device).view(wp.transformf) - return self._body_com_pose_b + self._body_com_pose_b_ta = None + if self._body_com_pose_b_ta is None: + self._body_com_pose_b_ta = ProxyArray(self._body_com_pose_b) + return self._body_com_pose_b_ta # Sliced properties (zero-copy pointer arithmetic on transformf) @property - def body_com_pos_w(self) -> wp.array: + def body_com_pos_w(self) -> ProxyArray: """Body CoM positions. Shape: (N, num_bodies), dtype=wp.vec3f.""" - t = self.body_com_pose_w - return wp.array(ptr=t.ptr, shape=t.shape, dtype=wp.vec3f, strides=t.strides, device=self.device) + t = self.body_com_pose_w.warp + return ProxyArray(wp.array(ptr=t.ptr, shape=t.shape, dtype=wp.vec3f, strides=t.strides, device=self.device)) @property - def body_com_quat_w(self) -> wp.array: + def body_com_quat_w(self) -> ProxyArray: """Body CoM orientations. Shape: (N, num_bodies), dtype=wp.quatf.""" - t = self.body_com_pose_w - return wp.array(ptr=t.ptr + 3 * 4, shape=t.shape, dtype=wp.quatf, strides=t.strides, device=self.device) + t = self.body_com_pose_w.warp + return ProxyArray( + wp.array(ptr=t.ptr + 3 * 4, shape=t.shape, dtype=wp.quatf, strides=t.strides, device=self.device) + ) @property - def body_com_lin_vel_w(self) -> wp.array: + def body_com_lin_vel_w(self) -> ProxyArray: """Body CoM linear velocities. Shape: (N, num_bodies), dtype=wp.vec3f.""" - v = self.body_com_vel_w - return wp.array(ptr=v.ptr, shape=v.shape, dtype=wp.vec3f, strides=v.strides, device=self.device) + v = self.body_com_vel_w.warp + return ProxyArray(wp.array(ptr=v.ptr, shape=v.shape, dtype=wp.vec3f, strides=v.strides, device=self.device)) @property - def body_com_ang_vel_w(self) -> wp.array: + def body_com_ang_vel_w(self) -> ProxyArray: """Body CoM angular velocities. Shape: (N, num_bodies), dtype=wp.vec3f.""" - v = self.body_com_vel_w - return wp.array(ptr=v.ptr + 3 * 4, shape=v.shape, dtype=wp.vec3f, strides=v.strides, device=self.device) + v = self.body_com_vel_w.warp + return ProxyArray( + wp.array(ptr=v.ptr + 3 * 4, shape=v.shape, dtype=wp.vec3f, strides=v.strides, device=self.device) + ) @property - def body_com_lin_acc_w(self) -> wp.array: + def body_com_lin_acc_w(self) -> ProxyArray: """Body CoM linear accelerations. Shape: (N, num_bodies), dtype=wp.vec3f.""" - v = self.body_com_acc_w - return wp.array(ptr=v.ptr, shape=v.shape, dtype=wp.vec3f, strides=v.strides, device=self.device) + v = self.body_com_acc_w.warp + return ProxyArray(wp.array(ptr=v.ptr, shape=v.shape, dtype=wp.vec3f, strides=v.strides, device=self.device)) @property - def body_com_ang_acc_w(self) -> wp.array: + def body_com_ang_acc_w(self) -> ProxyArray: """Body CoM angular accelerations. Shape: (N, num_bodies), dtype=wp.vec3f.""" - v = self.body_com_acc_w - return wp.array(ptr=v.ptr + 3 * 4, shape=v.shape, dtype=wp.vec3f, strides=v.strides, device=self.device) + v = self.body_com_acc_w.warp + return ProxyArray( + wp.array(ptr=v.ptr + 3 * 4, shape=v.shape, dtype=wp.vec3f, strides=v.strides, device=self.device) + ) @property - def body_com_pos_b(self) -> wp.array: + def body_com_pos_b(self) -> ProxyArray: """Body CoM positions in body frame. Shape: (N, num_bodies), dtype=wp.vec3f.""" - t = self.body_com_pose_b - return wp.array(ptr=t.ptr, shape=t.shape, dtype=wp.vec3f, strides=t.strides, device=self.device) + t = self.body_com_pose_b.warp + return ProxyArray(wp.array(ptr=t.ptr, shape=t.shape, dtype=wp.vec3f, strides=t.strides, device=self.device)) @property - def body_com_quat_b(self) -> wp.array: + def body_com_quat_b(self) -> ProxyArray: """Body CoM orientations in body frame. Shape: (N, num_bodies), dtype=wp.quatf.""" - t = self.body_com_pose_b - return wp.array(ptr=t.ptr + 3 * 4, shape=t.shape, dtype=wp.quatf, strides=t.strides, device=self.device) + t = self.body_com_pose_b.warp + return ProxyArray( + wp.array(ptr=t.ptr + 3 * 4, shape=t.shape, dtype=wp.quatf, strides=t.strides, device=self.device) + ) # -- Body properties -- @property - def body_mass(self) -> wp.array: + def body_mass(self) -> ProxyArray: """Body masses. Shape: (N, num_bodies).""" if self._body_mass is None: - return wp.ones((self._num_instances, self._num_bodies), dtype=wp.float32, device=self.device) - return self._body_mass + self._body_mass = wp.ones((self._num_instances, self._num_bodies), dtype=wp.float32, device=self.device) + self._body_mass_ta = None + if self._body_mass_ta is None: + self._body_mass_ta = ProxyArray(self._body_mass) + return self._body_mass_ta @property - def body_inertia(self) -> wp.array: + def body_inertia(self) -> ProxyArray: """Body inertias (flattened 3x3). Shape: (N, num_bodies, 9).""" if self._body_inertia is None: inertia_np = np.zeros((self._num_instances, self._num_bodies, 9), dtype=np.float32) inertia_np[..., 0] = 1.0 # Ixx inertia_np[..., 4] = 1.0 # Iyy inertia_np[..., 8] = 1.0 # Izz - return wp.array(inertia_np, dtype=wp.float32, device=self.device) - return self._body_inertia + self._body_inertia = wp.array(inertia_np, dtype=wp.float32, device=self.device) + self._body_inertia_ta = None + if self._body_inertia_ta is None: + self._body_inertia_ta = ProxyArray(self._body_inertia) + return self._body_inertia_ta # -- Derived properties -- @property - def projected_gravity_b(self) -> wp.array: + def projected_gravity_b(self) -> ProxyArray: """Gravity projection on bodies. Shape: (N, num_bodies), dtype=wp.vec3f.""" gravity_np = np.zeros((self._num_instances, self._num_bodies, 3), dtype=np.float32) gravity_np[..., 2] = -1.0 - return wp.array(gravity_np, dtype=wp.float32, device=self.device).view(wp.vec3f) + return ProxyArray(wp.array(gravity_np, dtype=wp.float32, device=self.device).view(wp.vec3f)) @property - def heading_w(self) -> wp.array: + def heading_w(self) -> ProxyArray: """Yaw heading per body. Shape: (N, num_bodies).""" - return wp.zeros((self._num_instances, self._num_bodies), dtype=wp.float32, device=self.device) + return ProxyArray(wp.zeros((self._num_instances, self._num_bodies), dtype=wp.float32, device=self.device)) @property - def body_link_lin_vel_b(self) -> wp.array: + def body_link_lin_vel_b(self) -> ProxyArray: """Body link linear velocities in body frame. Shape: (N, num_bodies), dtype=wp.vec3f.""" - return wp.clone(self.body_link_lin_vel_w, self.device) + return ProxyArray(wp.clone(self.body_link_lin_vel_w.warp, self.device)) @property - def body_link_ang_vel_b(self) -> wp.array: + def body_link_ang_vel_b(self) -> ProxyArray: """Body link angular velocities in body frame. Shape: (N, num_bodies), dtype=wp.vec3f.""" - return wp.clone(self.body_link_ang_vel_w, self.device) + return ProxyArray(wp.clone(self.body_link_ang_vel_w.warp, self.device)) @property - def body_com_lin_vel_b(self) -> wp.array: + def body_com_lin_vel_b(self) -> ProxyArray: """Body CoM linear velocities in body frame. Shape: (N, num_bodies), dtype=wp.vec3f.""" - return wp.clone(self.body_com_lin_vel_w, self.device) + return ProxyArray(wp.clone(self.body_com_lin_vel_w.warp, self.device)) @property - def body_com_ang_vel_b(self) -> wp.array: + def body_com_ang_vel_b(self) -> ProxyArray: """Body CoM angular velocities in body frame. Shape: (N, num_bodies), dtype=wp.vec3f.""" - return wp.clone(self.body_com_ang_vel_w, self.device) + return ProxyArray(wp.clone(self.body_com_ang_vel_w.warp, self.device)) # -- Body state (abstract) -- @property - def body_state_w(self) -> wp.array: + def body_state_w(self) -> ProxyArray: """Body states (link pose + CoM velocity). Shape: (N, num_bodies, 13).""" - pose = wp.to_torch(self.body_link_pose_w) - vel = wp.to_torch(self.body_com_vel_w) - return wp.from_torch(torch.cat([pose, vel], dim=-1)) + pose = self.body_link_pose_w.torch + vel = self.body_com_vel_w.torch + return ProxyArray(wp.from_torch(torch.cat([pose, vel], dim=-1))) # -- Shorthand properties (body_pose_w, body_pos_w, com_pos_b, etc.) -- # Inherited from BaseRigidObjectCollectionData @@ -336,33 +392,43 @@ def update(self, dt: float) -> None: def set_default_body_pose(self, value: torch.Tensor) -> None: self._default_body_pose = wp.from_torch(value.to(self.device).contiguous()).view(wp.transformf) + self._default_body_pose_ta = None def set_default_body_vel(self, value: torch.Tensor) -> None: self._default_body_vel = wp.from_torch(value.to(self.device).contiguous()).view(wp.spatial_vectorf) + self._default_body_vel_ta = None def set_body_link_pose_w(self, value: torch.Tensor) -> None: self._body_link_pose_w = wp.from_torch(value.to(self.device).contiguous()).view(wp.transformf) + self._body_link_pose_w_ta = None def set_body_link_vel_w(self, value: torch.Tensor) -> None: self._body_link_vel_w = wp.from_torch(value.to(self.device).contiguous()).view(wp.spatial_vectorf) + self._body_link_vel_w_ta = None def set_body_com_pose_w(self, value: torch.Tensor) -> None: self._body_com_pose_w = wp.from_torch(value.to(self.device).contiguous()).view(wp.transformf) + self._body_com_pose_w_ta = None def set_body_com_vel_w(self, value: torch.Tensor) -> None: self._body_com_vel_w = wp.from_torch(value.to(self.device).contiguous()).view(wp.spatial_vectorf) + self._body_com_vel_w_ta = None def set_body_com_acc_w(self, value: torch.Tensor) -> None: self._body_com_acc_w = wp.from_torch(value.to(self.device).contiguous()).view(wp.spatial_vectorf) + self._body_com_acc_w_ta = None def set_body_com_pose_b(self, value: torch.Tensor) -> None: self._body_com_pose_b = wp.from_torch(value.to(self.device).contiguous()).view(wp.transformf) + self._body_com_pose_b_ta = None def set_body_mass(self, value: torch.Tensor) -> None: self._body_mass = wp.from_torch(value.to(self.device).contiguous()) + self._body_mass_ta = None def set_body_inertia(self, value: torch.Tensor) -> None: self._body_inertia = wp.from_torch(value.to(self.device).contiguous()) + self._body_inertia_ta = None def set_mock_data(self, **kwargs) -> None: """Bulk setter for mock data.""" 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 index 82e134ccd6e9..4c6c9ad8c0b1 100644 --- a/source/isaaclab/isaaclab/test/mock_interfaces/sensors/mock_contact_sensor.py +++ b/source/isaaclab/isaaclab/test/mock_interfaces/sensors/mock_contact_sensor.py @@ -14,6 +14,8 @@ import torch import warp as wp +from isaaclab.utils.warp import ProxyArray + try: from isaaclab.sensors.contact_sensor.base_contact_sensor_data import BaseContactSensorData except (ImportError, ModuleNotFoundError): @@ -72,65 +74,97 @@ def __init__( self._last_contact_time: wp.array | None = None self._current_contact_time: wp.array | None = None + # ProxyArray caches + self._pos_w_ta: ProxyArray | None = None + self._quat_w_ta: ProxyArray | None = None + self._net_forces_w_ta: ProxyArray | None = None + self._net_forces_w_history_ta: ProxyArray | None = None + self._force_matrix_w_ta: ProxyArray | None = None + self._contact_pos_w_ta: ProxyArray | None = None + self._friction_forces_w_ta: ProxyArray | None = None + self._last_air_time_ta: ProxyArray | None = None + self._current_air_time_ta: ProxyArray | None = None + self._last_contact_time_ta: ProxyArray | None = None + self._current_contact_time_ta: ProxyArray | None = None + # -- Properties -- @property - def pos_w(self) -> wp.array | None: + def pos_w(self) -> ProxyArray | None: """Position of sensor origins in world frame. Shape: (N, B, 3).""" if self._pos_w is None: - return wp.zeros(shape=(self._num_instances, self._num_bodies, 3), dtype=wp.float32, device=self.device) - return self._pos_w + self._pos_w = wp.zeros( + shape=(self._num_instances, self._num_bodies, 3), dtype=wp.float32, device=self.device + ) + self._pos_w_ta = None + if self._pos_w_ta is None: + self._pos_w_ta = ProxyArray(self._pos_w) + return self._pos_w_ta @property - def quat_w(self) -> wp.array | None: + def quat_w(self) -> ProxyArray | None: """Orientation (w, x, y, z) in world frame. Shape: (N, B, 4).""" if self._quat_w is None: # Default to identity quaternion quat_np = np.zeros((self._num_instances, self._num_bodies, 4), dtype=np.float32) quat_np[..., 0] = 1.0 - return wp.array(quat_np, dtype=wp.float32, device=self.device) - return self._quat_w + self._quat_w = wp.array(quat_np, dtype=wp.float32, device=self.device) + self._quat_w_ta = None + if self._quat_w_ta is None: + self._quat_w_ta = ProxyArray(self._quat_w) + return self._quat_w_ta @property - def pose_w(self) -> wp.array | None: + def pose_w(self) -> ProxyArray | None: """Pose in world frame (pos + quat). Shape: (N, B, 7).""" - pos_t = wp.to_torch(self.pos_w) - quat_t = wp.to_torch(self.quat_w) + pos_t = self.pos_w.torch + quat_t = self.quat_w.torch pose_t = torch.cat([pos_t, quat_t], dim=-1) - return wp.from_torch(pose_t.contiguous(), dtype=wp.float32) + return ProxyArray(wp.from_torch(pose_t.contiguous(), dtype=wp.float32)) @property - def net_forces_w(self) -> wp.array: + def net_forces_w(self) -> ProxyArray: """Net normal contact forces in world frame. Shape: (N, B, 3).""" if self._net_forces_w is None: - return wp.zeros(shape=(self._num_instances, self._num_bodies, 3), dtype=wp.float32, device=self.device) - return self._net_forces_w + self._net_forces_w = wp.zeros( + shape=(self._num_instances, self._num_bodies, 3), dtype=wp.float32, device=self.device + ) + self._net_forces_w_ta = None + if self._net_forces_w_ta is None: + self._net_forces_w_ta = ProxyArray(self._net_forces_w) + return self._net_forces_w_ta @property - def net_forces_w_history(self) -> wp.array | None: + def net_forces_w_history(self) -> ProxyArray | 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 wp.zeros( + self._net_forces_w_history = wp.zeros( shape=(self._num_instances, self._history_length, self._num_bodies, 3), dtype=wp.float32, device=self.device, ) - return self._net_forces_w_history + self._net_forces_w_history_ta = None + if self._net_forces_w_history_ta is None: + self._net_forces_w_history_ta = ProxyArray(self._net_forces_w_history) + return self._net_forces_w_history_ta @property - def force_matrix_w(self) -> wp.array | None: + def force_matrix_w(self) -> ProxyArray | 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 wp.zeros( + self._force_matrix_w = wp.zeros( shape=(self._num_instances, self._num_bodies, self._num_filter_bodies, 3), dtype=wp.float32, device=self.device, ) - return self._force_matrix_w + self._force_matrix_w_ta = None + if self._force_matrix_w_ta is None: + self._force_matrix_w_ta = ProxyArray(self._force_matrix_w) + return self._force_matrix_w_ta @property def force_matrix_w_history(self) -> torch.Tensor | None: @@ -149,80 +183,111 @@ def force_matrix_w_history(self) -> torch.Tensor | None: return self._force_matrix_w_history @property - def contact_pos_w(self) -> wp.array | None: + def contact_pos_w(self) -> ProxyArray | 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 wp.zeros( + self._contact_pos_w = wp.zeros( shape=(self._num_instances, self._num_bodies, self._num_filter_bodies, 3), dtype=wp.float32, device=self.device, ) - return self._contact_pos_w + self._contact_pos_w_ta = None + if self._contact_pos_w_ta is None: + self._contact_pos_w_ta = ProxyArray(self._contact_pos_w) + return self._contact_pos_w_ta @property - def friction_forces_w(self) -> wp.array | None: + def friction_forces_w(self) -> ProxyArray | 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 wp.zeros( + self._friction_forces_w = wp.zeros( shape=(self._num_instances, self._num_bodies, self._num_filter_bodies, 3), dtype=wp.float32, device=self.device, ) - return self._friction_forces_w + self._friction_forces_w_ta = None + if self._friction_forces_w_ta is None: + self._friction_forces_w_ta = ProxyArray(self._friction_forces_w) + return self._friction_forces_w_ta @property - def last_air_time(self) -> wp.array: + def last_air_time(self) -> ProxyArray: """Time in air before last contact. Shape: (N, B).""" if self._last_air_time is None: - return wp.zeros(shape=(self._num_instances, self._num_bodies), dtype=wp.float32, device=self.device) - return self._last_air_time + self._last_air_time = wp.zeros( + shape=(self._num_instances, self._num_bodies), dtype=wp.float32, device=self.device + ) + self._last_air_time_ta = None + if self._last_air_time_ta is None: + self._last_air_time_ta = ProxyArray(self._last_air_time) + return self._last_air_time_ta @property - def current_air_time(self) -> wp.array: + def current_air_time(self) -> ProxyArray: """Current time in air. Shape: (N, B).""" if self._current_air_time is None: - return wp.zeros(shape=(self._num_instances, self._num_bodies), dtype=wp.float32, device=self.device) - return self._current_air_time + self._current_air_time = wp.zeros( + shape=(self._num_instances, self._num_bodies), dtype=wp.float32, device=self.device + ) + self._current_air_time_ta = None + if self._current_air_time_ta is None: + self._current_air_time_ta = ProxyArray(self._current_air_time) + return self._current_air_time_ta @property - def last_contact_time(self) -> wp.array: + def last_contact_time(self) -> ProxyArray: """Time in contact before last detach. Shape: (N, B).""" if self._last_contact_time is None: - return wp.zeros(shape=(self._num_instances, self._num_bodies), dtype=wp.float32, device=self.device) - return self._last_contact_time + self._last_contact_time = wp.zeros( + shape=(self._num_instances, self._num_bodies), dtype=wp.float32, device=self.device + ) + self._last_contact_time_ta = None + if self._last_contact_time_ta is None: + self._last_contact_time_ta = ProxyArray(self._last_contact_time) + return self._last_contact_time_ta @property - def current_contact_time(self) -> wp.array: + def current_contact_time(self) -> ProxyArray: """Current time in contact. Shape: (N, B).""" if self._current_contact_time is None: - return wp.zeros(shape=(self._num_instances, self._num_bodies), dtype=wp.float32, device=self.device) - return self._current_contact_time + self._current_contact_time = wp.zeros( + shape=(self._num_instances, self._num_bodies), dtype=wp.float32, device=self.device + ) + self._current_contact_time_ta = None + if self._current_contact_time_ta is None: + self._current_contact_time_ta = ProxyArray(self._current_contact_time) + return self._current_contact_time_ta # -- Setters -- def set_pos_w(self, value: torch.Tensor) -> None: """Set position in world frame.""" self._pos_w = wp.from_torch(value.to(self.device).contiguous(), dtype=wp.float32) + self._pos_w_ta = None def set_quat_w(self, value: torch.Tensor) -> None: """Set orientation in world frame.""" self._quat_w = wp.from_torch(value.to(self.device).contiguous(), dtype=wp.float32) + self._quat_w_ta = None def set_net_forces_w(self, value: torch.Tensor) -> None: """Set net contact forces.""" self._net_forces_w = wp.from_torch(value.to(self.device).contiguous(), dtype=wp.float32) + self._net_forces_w_ta = None def set_net_forces_w_history(self, value: torch.Tensor) -> None: """Set net forces history.""" self._net_forces_w_history = wp.from_torch(value.to(self.device).contiguous(), dtype=wp.float32) + self._net_forces_w_history_ta = None def set_force_matrix_w(self, value: torch.Tensor) -> None: """Set filtered contact forces.""" self._force_matrix_w = wp.from_torch(value.to(self.device).contiguous(), dtype=wp.float32) + self._force_matrix_w_ta = None def set_force_matrix_w_history(self, value: torch.Tensor) -> None: """Set filtered forces history. @@ -234,26 +299,32 @@ def set_force_matrix_w_history(self, value: torch.Tensor) -> None: def set_contact_pos_w(self, value: torch.Tensor) -> None: """Set contact point positions.""" self._contact_pos_w = wp.from_torch(value.to(self.device).contiguous(), dtype=wp.float32) + self._contact_pos_w_ta = None def set_friction_forces_w(self, value: torch.Tensor) -> None: """Set friction forces.""" self._friction_forces_w = wp.from_torch(value.to(self.device).contiguous(), dtype=wp.float32) + self._friction_forces_w_ta = None def set_last_air_time(self, value: torch.Tensor) -> None: """Set last air time.""" self._last_air_time = wp.from_torch(value.to(self.device).contiguous(), dtype=wp.float32) + self._last_air_time_ta = None def set_current_air_time(self, value: torch.Tensor) -> None: """Set current air time.""" self._current_air_time = wp.from_torch(value.to(self.device).contiguous(), dtype=wp.float32) + self._current_air_time_ta = None def set_last_contact_time(self, value: torch.Tensor) -> None: """Set last contact time.""" self._last_contact_time = wp.from_torch(value.to(self.device).contiguous(), dtype=wp.float32) + self._last_contact_time_ta = None def set_current_contact_time(self, value: torch.Tensor) -> None: """Set current contact time.""" self._current_contact_time = wp.from_torch(value.to(self.device).contiguous(), dtype=wp.float32) + self._current_contact_time_ta = None def set_mock_data( self, @@ -439,7 +510,7 @@ def compute_first_contact(self, dt: float, abs_tol: float = 1.0e-8) -> wp.array: Returns: Boolean warp array of shape (N, B) indicating first contact. """ - result = wp.to_torch(self._data.current_contact_time) < (dt + abs_tol) + result = self._data.current_contact_time.torch < (dt + abs_tol) return wp.from_torch(result) def compute_first_air(self, dt: float, abs_tol: float = 1.0e-8) -> wp.array: @@ -452,7 +523,7 @@ def compute_first_air(self, dt: float, abs_tol: float = 1.0e-8) -> wp.array: Returns: Boolean warp array of shape (N, B) indicating first air. """ - result = wp.to_torch(self._data.current_air_time) < (dt + abs_tol) + result = self._data.current_air_time.torch < (dt + abs_tol) return wp.from_torch(result) def reset(self, env_ids: Sequence[int] | None = None) -> None: 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 index e79f571642a2..5d4b9cad6153 100644 --- a/source/isaaclab/isaaclab/test/mock_interfaces/sensors/mock_frame_transformer.py +++ b/source/isaaclab/isaaclab/test/mock_interfaces/sensors/mock_frame_transformer.py @@ -14,6 +14,8 @@ import torch import warp as wp +from isaaclab.utils.warp import ProxyArray + try: from isaaclab.sensors.frame_transformer.base_frame_transformer_data import BaseFrameTransformerData except (ImportError, ModuleNotFoundError): @@ -63,6 +65,14 @@ def __init__( self._target_pos_source: wp.array | None = None self._target_quat_source: wp.array | None = None + # ProxyArray caches + self._source_pos_w_ta: ProxyArray | None = None + self._source_quat_w_ta: ProxyArray | None = None + self._target_pos_w_ta: ProxyArray | None = None + self._target_quat_w_ta: ProxyArray | None = None + self._target_pos_source_ta: ProxyArray | None = None + self._target_quat_source_ta: ProxyArray | None = None + # -- Properties -- @property @@ -71,106 +81,130 @@ def target_frame_names(self) -> list[str]: return self._target_frame_names @property - def source_pos_w(self) -> wp.array: + def source_pos_w(self) -> ProxyArray: """Position of source frame in world frame. Shape: (N, 3).""" if self._source_pos_w is None: - return wp.zeros(shape=(self._num_instances, 3), dtype=wp.float32, device=self.device) - return self._source_pos_w + self._source_pos_w = wp.zeros(shape=(self._num_instances, 3), dtype=wp.float32, device=self.device) + self._source_pos_w_ta = None + if self._source_pos_w_ta is None: + self._source_pos_w_ta = ProxyArray(self._source_pos_w) + return self._source_pos_w_ta @property - def source_quat_w(self) -> wp.array: + def source_quat_w(self) -> ProxyArray: """Orientation (w, x, y, z) of source frame in world frame. Shape: (N, 4).""" if self._source_quat_w is None: quat_np = np.zeros((self._num_instances, 4), dtype=np.float32) quat_np[:, 0] = 1.0 - return wp.array(quat_np, dtype=wp.float32, device=self.device) - return self._source_quat_w + self._source_quat_w = wp.array(quat_np, dtype=wp.float32, device=self.device) + self._source_quat_w_ta = None + if self._source_quat_w_ta is None: + self._source_quat_w_ta = ProxyArray(self._source_quat_w) + return self._source_quat_w_ta @property - def source_pose_w(self) -> wp.array: + def source_pose_w(self) -> ProxyArray: """Pose of source frame in world frame. Shape: (N, 7).""" - pos_t = wp.to_torch(self.source_pos_w) - quat_t = wp.to_torch(self.source_quat_w) + pos_t = self.source_pos_w.torch + quat_t = self.source_quat_w.torch pose_t = torch.cat([pos_t, quat_t], dim=-1) - return wp.from_torch(pose_t.contiguous(), dtype=wp.float32) + return ProxyArray(wp.from_torch(pose_t.contiguous(), dtype=wp.float32)) @property - def target_pos_w(self) -> wp.array: + def target_pos_w(self) -> ProxyArray: """Position of target frames in world frame. Shape: (N, M, 3).""" if self._target_pos_w is None: - return wp.zeros( + self._target_pos_w = wp.zeros( shape=(self._num_instances, self._num_target_frames, 3), dtype=wp.float32, device=self.device ) - return self._target_pos_w + self._target_pos_w_ta = None + if self._target_pos_w_ta is None: + self._target_pos_w_ta = ProxyArray(self._target_pos_w) + return self._target_pos_w_ta @property - def target_quat_w(self) -> wp.array: + def target_quat_w(self) -> ProxyArray: """Orientation (w, x, y, z) of target frames in world frame. Shape: (N, M, 4).""" if self._target_quat_w is None: quat_np = np.zeros((self._num_instances, self._num_target_frames, 4), dtype=np.float32) quat_np[..., 0] = 1.0 - return wp.array(quat_np, dtype=wp.float32, device=self.device) - return self._target_quat_w + self._target_quat_w = wp.array(quat_np, dtype=wp.float32, device=self.device) + self._target_quat_w_ta = None + if self._target_quat_w_ta is None: + self._target_quat_w_ta = ProxyArray(self._target_quat_w) + return self._target_quat_w_ta @property - def target_pose_w(self) -> wp.array: + def target_pose_w(self) -> ProxyArray: """Pose of target frames in world frame. Shape: (N, M, 7).""" - pos_t = wp.to_torch(self.target_pos_w) - quat_t = wp.to_torch(self.target_quat_w) + pos_t = self.target_pos_w.torch + quat_t = self.target_quat_w.torch pose_t = torch.cat([pos_t, quat_t], dim=-1) - return wp.from_torch(pose_t.contiguous(), dtype=wp.float32) + return ProxyArray(wp.from_torch(pose_t.contiguous(), dtype=wp.float32)) @property - def target_pos_source(self) -> wp.array: + def target_pos_source(self) -> ProxyArray: """Position of target frames relative to source frame. Shape: (N, M, 3).""" if self._target_pos_source is None: - return wp.zeros( + self._target_pos_source = wp.zeros( shape=(self._num_instances, self._num_target_frames, 3), dtype=wp.float32, device=self.device ) - return self._target_pos_source + self._target_pos_source_ta = None + if self._target_pos_source_ta is None: + self._target_pos_source_ta = ProxyArray(self._target_pos_source) + return self._target_pos_source_ta @property - def target_quat_source(self) -> wp.array: + def target_quat_source(self) -> ProxyArray: """Orientation (w, x, y, z) of target frames relative to source frame. Shape: (N, M, 4).""" if self._target_quat_source is None: quat_np = np.zeros((self._num_instances, self._num_target_frames, 4), dtype=np.float32) quat_np[..., 0] = 1.0 - return wp.array(quat_np, dtype=wp.float32, device=self.device) - return self._target_quat_source + self._target_quat_source = wp.array(quat_np, dtype=wp.float32, device=self.device) + self._target_quat_source_ta = None + if self._target_quat_source_ta is None: + self._target_quat_source_ta = ProxyArray(self._target_quat_source) + return self._target_quat_source_ta @property - def target_pose_source(self) -> wp.array: + def target_pose_source(self) -> ProxyArray: """Pose of target frames relative to source frame. Shape: (N, M, 7).""" - pos_t = wp.to_torch(self.target_pos_source) - quat_t = wp.to_torch(self.target_quat_source) + pos_t = self.target_pos_source.torch + quat_t = self.target_quat_source.torch pose_t = torch.cat([pos_t, quat_t], dim=-1) - return wp.from_torch(pose_t.contiguous(), dtype=wp.float32) + return ProxyArray(wp.from_torch(pose_t.contiguous(), dtype=wp.float32)) # -- Setters -- def set_source_pos_w(self, value: torch.Tensor) -> None: """Set source position in world frame.""" self._source_pos_w = wp.from_torch(value.to(self.device).contiguous(), dtype=wp.float32) + self._source_pos_w_ta = None def set_source_quat_w(self, value: torch.Tensor) -> None: """Set source orientation in world frame.""" self._source_quat_w = wp.from_torch(value.to(self.device).contiguous(), dtype=wp.float32) + self._source_quat_w_ta = None def set_target_pos_w(self, value: torch.Tensor) -> None: """Set target positions in world frame.""" self._target_pos_w = wp.from_torch(value.to(self.device).contiguous(), dtype=wp.float32) + self._target_pos_w_ta = None def set_target_quat_w(self, value: torch.Tensor) -> None: """Set target orientations in world frame.""" self._target_quat_w = wp.from_torch(value.to(self.device).contiguous(), dtype=wp.float32) + self._target_quat_w_ta = None def set_target_pos_source(self, value: torch.Tensor) -> None: """Set target positions relative to source.""" self._target_pos_source = wp.from_torch(value.to(self.device).contiguous(), dtype=wp.float32) + self._target_pos_source_ta = None def set_target_quat_source(self, value: torch.Tensor) -> None: """Set target orientations relative to source.""" self._target_quat_source = wp.from_torch(value.to(self.device).contiguous(), dtype=wp.float32) + self._target_quat_source_ta = None def set_mock_data( self, diff --git a/source/isaaclab/isaaclab/test/mock_interfaces/sensors/mock_imu.py b/source/isaaclab/isaaclab/test/mock_interfaces/sensors/mock_imu.py index 7fb6dac22978..f32f0a5f567a 100644 --- a/source/isaaclab/isaaclab/test/mock_interfaces/sensors/mock_imu.py +++ b/source/isaaclab/isaaclab/test/mock_interfaces/sensors/mock_imu.py @@ -12,6 +12,8 @@ import torch import warp as wp +from isaaclab.utils.warp import ProxyArray + try: from isaaclab.sensors.imu.base_imu_data import BaseImuData except (ImportError, ModuleNotFoundError): @@ -46,31 +48,43 @@ def __init__(self, num_instances: int, device: str = "cpu"): self._ang_vel_b: wp.array | None = None self._lin_acc_b: wp.array | None = None + # ProxyArray caches + self._ang_vel_b_ta: ProxyArray | None = None + self._lin_acc_b_ta: ProxyArray | None = None + # -- Properties -- @property - def ang_vel_b(self) -> wp.array: + def ang_vel_b(self) -> ProxyArray: """Angular velocity in IMU body frame [rad/s]. Shape: (N, 3).""" if self._ang_vel_b is None: - return wp.zeros(shape=(self._num_instances, 3), dtype=wp.float32, device=self.device) - return self._ang_vel_b + self._ang_vel_b = wp.zeros(shape=(self._num_instances, 3), dtype=wp.float32, device=self.device) + self._ang_vel_b_ta = None + if self._ang_vel_b_ta is None: + self._ang_vel_b_ta = ProxyArray(self._ang_vel_b) + return self._ang_vel_b_ta @property - def lin_acc_b(self) -> wp.array: + def lin_acc_b(self) -> ProxyArray: """Linear acceleration in IMU body frame [m/s^2]. Shape: (N, 3).""" if self._lin_acc_b is None: - return wp.zeros(shape=(self._num_instances, 3), dtype=wp.float32, device=self.device) - return self._lin_acc_b + self._lin_acc_b = wp.zeros(shape=(self._num_instances, 3), dtype=wp.float32, device=self.device) + self._lin_acc_b_ta = None + if self._lin_acc_b_ta is None: + self._lin_acc_b_ta = ProxyArray(self._lin_acc_b) + return self._lin_acc_b_ta # -- Setters -- def set_ang_vel_b(self, value: torch.Tensor) -> None: """Set angular velocity in body frame.""" self._ang_vel_b = wp.from_torch(value.to(self.device).contiguous(), dtype=wp.float32) + self._ang_vel_b_ta = None def set_lin_acc_b(self, value: torch.Tensor) -> None: """Set linear acceleration in body frame.""" self._lin_acc_b = wp.from_torch(value.to(self.device).contiguous(), dtype=wp.float32) + self._lin_acc_b_ta = None def set_mock_data( self, diff --git a/source/isaaclab/isaaclab/test/mock_interfaces/sensors/mock_pva.py b/source/isaaclab/isaaclab/test/mock_interfaces/sensors/mock_pva.py index 2f34b6fb2ff1..0fb24b551b1f 100644 --- a/source/isaaclab/isaaclab/test/mock_interfaces/sensors/mock_pva.py +++ b/source/isaaclab/isaaclab/test/mock_interfaces/sensors/mock_pva.py @@ -13,6 +13,8 @@ import torch import warp as wp +from isaaclab.utils.warp import ProxyArray + try: from isaaclab.sensors.pva.base_pva_data import BasePvaData except (ImportError, ModuleNotFoundError): @@ -53,100 +55,137 @@ def __init__(self, num_instances: int, device: str = "cpu"): self._lin_acc_b: wp.array | None = None self._ang_acc_b: wp.array | None = None + # ProxyArray caches + self._pos_w_ta: ProxyArray | None = None + self._quat_w_ta: ProxyArray | None = None + self._projected_gravity_b_ta: ProxyArray | None = None + self._lin_vel_b_ta: ProxyArray | None = None + self._ang_vel_b_ta: ProxyArray | None = None + self._lin_acc_b_ta: ProxyArray | None = None + self._ang_acc_b_ta: ProxyArray | None = None + # -- Properties -- @property - def pos_w(self) -> wp.array: + def pos_w(self) -> ProxyArray: """Position of sensor origin in world frame. Shape: (N, 3).""" if self._pos_w is None: - return wp.zeros(shape=(self._num_instances, 3), dtype=wp.float32, device=self.device) - return self._pos_w + self._pos_w = wp.zeros(shape=(self._num_instances, 3), dtype=wp.float32, device=self.device) + self._pos_w_ta = None + if self._pos_w_ta is None: + self._pos_w_ta = ProxyArray(self._pos_w) + return self._pos_w_ta @property - def quat_w(self) -> wp.array: + def quat_w(self) -> ProxyArray: """Orientation (x, y, z, w) in world frame. Shape: (N, 4).""" if self._quat_w is None: # Default to identity quaternion (x, y, z, w) = (0, 0, 0, 1) quat_np = np.zeros((self._num_instances, 4), dtype=np.float32) quat_np[:, 3] = 1.0 # w component is at index 3 in XYZW format - return wp.array(quat_np, dtype=wp.float32, device=self.device) - return self._quat_w + self._quat_w = wp.array(quat_np, dtype=wp.float32, device=self.device) + self._quat_w_ta = None + if self._quat_w_ta is None: + self._quat_w_ta = ProxyArray(self._quat_w) + return self._quat_w_ta @property - def pose_w(self) -> wp.array: + def pose_w(self) -> ProxyArray: """Pose in world frame (pos + quat). Shape: (N, 7).""" - pos_t = wp.to_torch(self.pos_w) - quat_t = wp.to_torch(self.quat_w) + pos_t = self.pos_w.torch + quat_t = self.quat_w.torch pose_t = torch.cat([pos_t, quat_t], dim=-1) - return wp.from_torch(pose_t.contiguous(), dtype=wp.float32) + return ProxyArray(wp.from_torch(pose_t.contiguous(), dtype=wp.float32)) @property - def projected_gravity_b(self) -> wp.array: + def projected_gravity_b(self) -> ProxyArray: """Gravity direction in PVA body frame. Shape: (N, 3).""" if self._projected_gravity_b is None: # Default gravity pointing down in body frame gravity_np = np.zeros((self._num_instances, 3), dtype=np.float32) gravity_np[:, 2] = -1.0 - return wp.array(gravity_np, dtype=wp.float32, device=self.device) - return self._projected_gravity_b + self._projected_gravity_b = wp.array(gravity_np, dtype=wp.float32, device=self.device) + self._projected_gravity_b_ta = None + if self._projected_gravity_b_ta is None: + self._projected_gravity_b_ta = ProxyArray(self._projected_gravity_b) + return self._projected_gravity_b_ta @property - def lin_vel_b(self) -> wp.array: + def lin_vel_b(self) -> ProxyArray: """Linear velocity in PVA body frame. Shape: (N, 3).""" if self._lin_vel_b is None: - return wp.zeros(shape=(self._num_instances, 3), dtype=wp.float32, device=self.device) - return self._lin_vel_b + self._lin_vel_b = wp.zeros(shape=(self._num_instances, 3), dtype=wp.float32, device=self.device) + self._lin_vel_b_ta = None + if self._lin_vel_b_ta is None: + self._lin_vel_b_ta = ProxyArray(self._lin_vel_b) + return self._lin_vel_b_ta @property - def ang_vel_b(self) -> wp.array: + def ang_vel_b(self) -> ProxyArray: """Angular velocity in PVA body frame. Shape: (N, 3).""" if self._ang_vel_b is None: - return wp.zeros(shape=(self._num_instances, 3), dtype=wp.float32, device=self.device) - return self._ang_vel_b + self._ang_vel_b = wp.zeros(shape=(self._num_instances, 3), dtype=wp.float32, device=self.device) + self._ang_vel_b_ta = None + if self._ang_vel_b_ta is None: + self._ang_vel_b_ta = ProxyArray(self._ang_vel_b) + return self._ang_vel_b_ta @property - def lin_acc_b(self) -> wp.array: + def lin_acc_b(self) -> ProxyArray: """Linear acceleration in PVA body frame. Shape: (N, 3).""" if self._lin_acc_b is None: - return wp.zeros(shape=(self._num_instances, 3), dtype=wp.float32, device=self.device) - return self._lin_acc_b + self._lin_acc_b = wp.zeros(shape=(self._num_instances, 3), dtype=wp.float32, device=self.device) + self._lin_acc_b_ta = None + if self._lin_acc_b_ta is None: + self._lin_acc_b_ta = ProxyArray(self._lin_acc_b) + return self._lin_acc_b_ta @property - def ang_acc_b(self) -> wp.array: + def ang_acc_b(self) -> ProxyArray: """Angular acceleration in PVA body frame. Shape: (N, 3).""" if self._ang_acc_b is None: - return wp.zeros(shape=(self._num_instances, 3), dtype=wp.float32, device=self.device) - return self._ang_acc_b + self._ang_acc_b = wp.zeros(shape=(self._num_instances, 3), dtype=wp.float32, device=self.device) + self._ang_acc_b_ta = None + if self._ang_acc_b_ta is None: + self._ang_acc_b_ta = ProxyArray(self._ang_acc_b) + return self._ang_acc_b_ta # -- Setters -- def set_pos_w(self, value: torch.Tensor) -> None: """Set position in world frame.""" self._pos_w = wp.from_torch(value.to(self.device).contiguous(), dtype=wp.float32) + self._pos_w_ta = None def set_quat_w(self, value: torch.Tensor) -> None: """Set orientation in world frame.""" self._quat_w = wp.from_torch(value.to(self.device).contiguous(), dtype=wp.float32) + self._quat_w_ta = None def set_projected_gravity_b(self, value: torch.Tensor) -> None: """Set projected gravity in body frame.""" self._projected_gravity_b = wp.from_torch(value.to(self.device).contiguous(), dtype=wp.float32) + self._projected_gravity_b_ta = None def set_lin_vel_b(self, value: torch.Tensor) -> None: """Set linear velocity in body frame.""" self._lin_vel_b = wp.from_torch(value.to(self.device).contiguous(), dtype=wp.float32) + self._lin_vel_b_ta = None def set_ang_vel_b(self, value: torch.Tensor) -> None: """Set angular velocity in body frame.""" self._ang_vel_b = wp.from_torch(value.to(self.device).contiguous(), dtype=wp.float32) + self._ang_vel_b_ta = None def set_lin_acc_b(self, value: torch.Tensor) -> None: """Set linear acceleration in body frame.""" self._lin_acc_b = wp.from_torch(value.to(self.device).contiguous(), dtype=wp.float32) + self._lin_acc_b_ta = None def set_ang_acc_b(self, value: torch.Tensor) -> None: """Set angular acceleration in body frame.""" self._ang_acc_b = wp.from_torch(value.to(self.device).contiguous(), dtype=wp.float32) + self._ang_acc_b_ta = None def set_mock_data( self, diff --git a/source/isaaclab/isaaclab/utils/warp/__init__.py b/source/isaaclab/isaaclab/utils/warp/__init__.py index e5a4264c81c7..569b80c9b089 100644 --- a/source/isaaclab/isaaclab/utils/warp/__init__.py +++ b/source/isaaclab/isaaclab/utils/warp/__init__.py @@ -5,6 +5,8 @@ """Sub-module containing operations based on warp.""" +import warnings + import warp as wp wp.config.quiet = True @@ -13,3 +15,47 @@ from isaaclab.utils.module import lazy_export lazy_export() + +# Avoid a circular import at module load: `.proxy_array` imports warp, which is +# already loaded above. Importing it here ensures the class is available inside +# the shim. +from .proxy_array import ProxyArray # noqa: E402 + +_WP_TO_TORCH_ORIGINAL = wp.to_torch +_WP_TO_TORCH_WARNED = False + + +def _wp_to_torch_with_proxyarray(a, requires_grad=None): + """Shim for :func:`warp.to_torch` that gracefully handles :class:`ProxyArray`. + + Without this shim, ``wp.to_torch(proxy)`` would fail with + ``AttributeError: 'ProxyArray' object has no attribute 'requires_grad'`` + because :class:`ProxyArray` intentionally doesn't replicate the full + ``wp.array`` attribute surface. Users and third-party code that still use + ``wp.to_torch(asset.data.)`` from before the ProxyArray migration + would break hard. + + The shim routes :class:`ProxyArray` arguments to their cached ``.torch`` + view (a zero-copy :class:`torch.Tensor` of the same underlying memory) and + emits a one-shot :class:`DeprecationWarning`. For any other input type, + the original :func:`warp.to_torch` handles the call as before. + """ + global _WP_TO_TORCH_WARNED + if isinstance(a, ProxyArray): + if not _WP_TO_TORCH_WARNED: + _WP_TO_TORCH_WARNED = True + warnings.warn( + "wp.to_torch() is deprecated; use the `.torch` accessor on" + " the ProxyArray directly (e.g. `asset.data.joint_pos.torch`).", + DeprecationWarning, + stacklevel=2, + ) + return a.torch + return _WP_TO_TORCH_ORIGINAL(a, requires_grad=requires_grad) + + +# Patch at both the top-level ``warp`` namespace and the underlying module so +# callers using ``import warp as wp`` and rare ``from warp._src.torch import +# to_torch`` patterns both pick up the shim. +wp.to_torch = _wp_to_torch_with_proxyarray +wp._src.torch.to_torch = _wp_to_torch_with_proxyarray # noqa: SLF001 diff --git a/source/isaaclab/isaaclab/utils/warp/__init__.pyi b/source/isaaclab/isaaclab/utils/warp/__init__.pyi index ea25f04435d3..ab1de52f39f7 100644 --- a/source/isaaclab/isaaclab/utils/warp/__init__.pyi +++ b/source/isaaclab/isaaclab/utils/warp/__init__.pyi @@ -4,6 +4,7 @@ # SPDX-License-Identifier: BSD-3-Clause __all__ = [ + "ProxyArray", "convert_to_warp_mesh", "raycast_dynamic_meshes", "raycast_mesh", @@ -11,3 +12,4 @@ __all__ = [ ] from .ops import convert_to_warp_mesh, raycast_dynamic_meshes, raycast_mesh, raycast_single_mesh +from .proxy_array import ProxyArray diff --git a/source/isaaclab/isaaclab/utils/warp/proxy_array.py b/source/isaaclab/isaaclab/utils/warp/proxy_array.py new file mode 100644 index 000000000000..912e44798877 --- /dev/null +++ b/source/isaaclab/isaaclab/utils/warp/proxy_array.py @@ -0,0 +1,338 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Warp-first dual-access array wrapper with explicit ``.torch`` and ``.warp`` accessors. + +Inspired by ProxyArray from mujocolab/mjlab (BSD-3-Clause). +""" + +from __future__ import annotations + +import os +import warnings +from typing import ClassVar + +import torch +import warp as wp + +_QUATF_ACCESS_WARN_ENV = "WARN_ON_TORCH_QUATF_ACCESS" +"""Environment variable that, when set to ``"1"``, makes :attr:`ProxyArray.torch` +emit a :class:`UserWarning` on every read of a ``wp.quatf``-typed array. Used as a +runtime aid for tracking down call sites that may still assume Isaac Lab 2.x's +``(w, x, y, z)`` quaternion convention after the migration to Isaac Lab 3.x's +``(x, y, z, w)`` convention. See the Isaac Lab 3.0 migration guide for details.""" + + +class ProxyArray: + """Warp-first array wrapper providing cached zero-copy ``.torch`` and ``.warp`` accessors. + + This class wraps a :class:`warp.array` and provides: + + * A ``.warp`` property that returns the original warp array (for kernel interop). + * A ``.torch`` property that returns a cached, zero-copy :class:`torch.Tensor` view + (via :func:`warp.to_torch`). + * Convenience properties (``shape``, ``dtype``, ``device``) delegated to the warp array. + * A deprecation bridge for common torch functions, indexing, and arithmetic/comparison + operators while emitting a one-time :class:`DeprecationWarning`. Tensor instance methods + such as ``clone()`` are not forwarded; use explicit ``.torch`` access for those. + + Example: + + .. code-block:: python + + import warp as wp + from isaaclab.utils.warp.proxy_array import ProxyArray + + arr = wp.zeros(100, dtype=wp.vec3f, device="cuda:0") + ta = ProxyArray(arr) + + # Explicit access (preferred) + ta.warp # -> wp.array, shape (100,), dtype vec3f + ta.torch # -> torch.Tensor, shape (100, 3) + + # Deprecation bridge (warns once, then silent) + result = ta + 1.0 # works, emits DeprecationWarning + """ + + _deprecation_warned: ClassVar[bool] = False + """Class-level flag ensuring the deprecation warning is emitted at most once.""" + + def __init__(self, wp_array: wp.array) -> None: + """Initialize the ProxyArray wrapper. + + The instance is immutable after construction: the wrapped ``wp.array`` cannot + be reassigned. If the underlying simulation memory is re-allocated, construct + a new :class:`ProxyArray` instead of mutating an existing one. + + Args: + wp_array: The warp array to wrap. + + Raises: + TypeError: If ``wp_array`` is not a :class:`warp.array`. + """ + if not isinstance(wp_array, wp.array): + raise TypeError( + f"ProxyArray expects a warp.array, got {type(wp_array).__name__}." + " If you have a ProxyArray, use it directly instead of wrapping it again." + ) + # Bypass __setattr__ for the two internal fields — everything else raises. + object.__setattr__(self, "_warp", wp_array) + object.__setattr__(self, "_torch_cache", None) + # Cached once at construction so the .torch read path stays a constant-time + # check; only used when the WARN_ON_TORCH_QUATF_ACCESS env var is set. + object.__setattr__(self, "_is_quatf", wp_array.dtype is wp.quatf) + + def __setattr__(self, name: str, value) -> None: + """Forbid mutation of ProxyArray instances except for the internal torch cache. + + The torch view is populated lazily on first ``.torch`` access; that is the + only allowed post-init state change. Every other write raises + :class:`AttributeError` so callers don't accidentally re-point the wrapper. + """ + if name == "_torch_cache": + object.__setattr__(self, name, value) + return + raise AttributeError( + f"ProxyArray is immutable; cannot set attribute {name!r}." + " Construct a new ProxyArray instead of mutating an existing one." + ) + + @staticmethod + def _quatf_access_warning_enabled() -> bool: + """Return ``True`` when the ``WARN_ON_TORCH_QUATF_ACCESS`` env var is set to ``"1"``. + + Read on every :attr:`torch` access to keep the flag dynamic — a single + ``os.environ`` lookup is cheap relative to the warp/torch interop work + that follows. + """ + return os.environ.get(_QUATF_ACCESS_WARN_ENV, "0") == "1" + + # ------------------------------------------------------------------ + # Core accessors + # ------------------------------------------------------------------ + + @property + def warp(self) -> wp.array: + """The underlying warp array.""" + return self._warp + + @property + def torch(self) -> torch.Tensor: + """A cached, zero-copy :class:`torch.Tensor` view of the warp array. + + The tensor is created on first access via :func:`warp.to_torch` and cached + for subsequent calls. Since this is a zero-copy view, modifications to the + tensor are visible through the warp array and vice versa. + + When the underlying warp array has dtype ``wp.quatf`` and the + ``WARN_ON_TORCH_QUATF_ACCESS`` environment variable is set to ``"1"``, + each read emits a :class:`UserWarning` pointing at the call site. This + is a runtime aid for migrating Isaac Lab 2.x code (which used the + ``(w, x, y, z)`` quaternion convention) to Isaac Lab 3.x's + ``(x, y, z, w)`` convention. + """ + if self._is_quatf and self._quatf_access_warning_enabled(): + warnings.warn( + "Reading .torch on a wp.quatf-typed ProxyArray. The Isaac Lab" + " quaternion convention changed from (w, x, y, z) in 2.x to" + " (x, y, z, w) in 3.x. If your code assumes the old order," + " this is likely the source of incorrect rotations." + f" Unset {_QUATF_ACCESS_WARN_ENV} to silence this warning.", + UserWarning, + stacklevel=2, + ) + if self._torch_cache is None: + self._torch_cache = wp.to_torch(self._warp) + return self._torch_cache + + # ------------------------------------------------------------------ + # Convenience properties + # ------------------------------------------------------------------ + + @property + def shape(self) -> tuple[int, ...]: + """Shape of the underlying warp array.""" + return self._warp.shape + + @property + def dtype(self): + """Warp dtype of the underlying array.""" + return self._warp.dtype + + @property + def device(self) -> str: + """Device string of the underlying warp array.""" + return self._warp.device + + def __len__(self) -> int: + """Return the size of the first dimension.""" + return self._warp.shape[0] + + def __repr__(self) -> str: + """Return a string representation of the ProxyArray.""" + return f"ProxyArray(shape={self.shape}, dtype={self.dtype}, device={self.device})" + + # ------------------------------------------------------------------ + # Warp kernel interop + # ------------------------------------------------------------------ + + @property + def __cuda_array_interface__(self): + """Delegate the CUDA array interface to the underlying warp array. + + This allows a ``ProxyArray`` to be passed directly as an argument to + :func:`warp.launch` without explicitly accessing ``.warp``. + + Raises: + AttributeError: If the underlying warp array is not on a CUDA device. + """ + return self._warp.__cuda_array_interface__ + + @property + def __array_interface__(self): + """Delegate the NumPy array interface to the underlying warp array. + + This allows a ``ProxyArray`` to be passed directly as an argument to + :func:`warp.launch` on CPU without explicitly accessing ``.warp``. + + Raises: + AttributeError: If the underlying warp array is not on a CPU device. + """ + return self._warp.__array_interface__ + + # ------------------------------------------------------------------ + # Indexing (deprecation bridge — delegates to .torch) + # ------------------------------------------------------------------ + + def __getitem__(self, key): + """Index into the torch view of this array. + + Supports all torch indexing: ``int``, ``slice``, ``tuple``, + boolean masks, and fancy indexing (multi-dimensional). + """ + self._warn_implicit() + return self.torch[key] + + def __setitem__(self, key, value): + """Write through the torch view into the shared warp memory. + + Supports all torch indexing: ``int``, ``slice``, ``tuple``, + boolean masks, and fancy indexing (multi-dimensional). + """ + self._warn_implicit() + self.torch[key] = value + + # ------------------------------------------------------------------ + # Deprecation bridge + # ------------------------------------------------------------------ + + @classmethod + def _warn_implicit(cls) -> None: + """Emit a one-time deprecation warning for implicit torch usage.""" + if not cls._deprecation_warned: + cls._deprecation_warned = True + warnings.warn( + "Implicit use of ProxyArray as a torch.Tensor is deprecated. " + "Use the explicit .torch property instead (e.g., array.torch).", + DeprecationWarning, + stacklevel=3, + ) + + @classmethod + def __torch_function__(cls, func, types, args=(), kwargs=None): + """Enable torch operations on ProxyArray by unwrapping to ``.torch``. + + This method is called by PyTorch when a torch function receives a + ``ProxyArray`` as an argument. It unwraps all ``ProxyArray`` instances + to their ``.torch`` tensors and delegates to the original function. + """ + if kwargs is None: + kwargs = {} + cls._warn_implicit() + + def unwrap(x): + if isinstance(x, ProxyArray): + return x.torch + if isinstance(x, (list, tuple)): + return type(x)(unwrap(i) for i in x) + return x + + args = unwrap(args) + kwargs = {k: unwrap(v) for k, v in kwargs.items()} + return func(*args, **kwargs) + + # ------------------------------------------------------------------ + # Arithmetic operators + # ------------------------------------------------------------------ + + def _binop(self, other, op: str) -> torch.Tensor: + """Helper for binary and reflected binary operations.""" + self._warn_implicit() + other_val = other.torch if isinstance(other, ProxyArray) else other + return getattr(self.torch, op)(other_val) + + def __add__(self, other) -> torch.Tensor: + return self._binop(other, "__add__") + + def __radd__(self, other) -> torch.Tensor: + return self._binop(other, "__radd__") + + def __sub__(self, other) -> torch.Tensor: + return self._binop(other, "__sub__") + + def __rsub__(self, other) -> torch.Tensor: + return self._binop(other, "__rsub__") + + def __mul__(self, other) -> torch.Tensor: + return self._binop(other, "__mul__") + + def __rmul__(self, other) -> torch.Tensor: + return self._binop(other, "__rmul__") + + def __truediv__(self, other) -> torch.Tensor: + return self._binop(other, "__truediv__") + + def __rtruediv__(self, other) -> torch.Tensor: + return self._binop(other, "__rtruediv__") + + def __pow__(self, other) -> torch.Tensor: + return self._binop(other, "__pow__") + + def __rpow__(self, other) -> torch.Tensor: + return self._binop(other, "__rpow__") + + def __neg__(self) -> torch.Tensor: + self._warn_implicit() + return -self.torch + + def __pos__(self) -> torch.Tensor: + self._warn_implicit() + return +self.torch + + def __abs__(self) -> torch.Tensor: + self._warn_implicit() + return abs(self.torch) + + # ------------------------------------------------------------------ + # Comparison operators + # ------------------------------------------------------------------ + + def __eq__(self, other) -> torch.Tensor: + return self._binop(other, "__eq__") + + def __ne__(self, other) -> torch.Tensor: + return self._binop(other, "__ne__") + + def __lt__(self, other) -> torch.Tensor: + return self._binop(other, "__lt__") + + def __le__(self, other) -> torch.Tensor: + return self._binop(other, "__le__") + + def __gt__(self, other) -> torch.Tensor: + return self._binop(other, "__gt__") + + def __ge__(self, other) -> torch.Tensor: + return self._binop(other, "__ge__") diff --git a/source/isaaclab/isaaclab/utils/wrench_composer.py b/source/isaaclab/isaaclab/utils/wrench_composer.py index e348697306d5..6c9390e97d7b 100644 --- a/source/isaaclab/isaaclab/utils/wrench_composer.py +++ b/source/isaaclab/isaaclab/utils/wrench_composer.py @@ -13,6 +13,7 @@ import torch import warp as wp +from isaaclab.utils.warp import ProxyArray from isaaclab.utils.warp.kernels import ( add_forces_to_dual_buffers_index, add_forces_to_dual_buffers_mask, @@ -66,11 +67,11 @@ def __init__(self, asset: BaseArticulation | BaseRigidObject | BaseRigidObjectCo self._active = False self._dirty = False if hasattr(self._asset.data, "body_com_pos_w"): - self._get_com_pos_fn = lambda a=self._asset: a.data.body_com_pos_w + self._get_com_pos_fn = lambda a=self._asset: a.data.body_com_pos_w.warp else: raise ValueError(f"Unsupported asset type: {self._asset.__class__.__name__}") if hasattr(self._asset.data, "body_link_quat_w"): - self._get_link_quat_fn = lambda a=self._asset: a.data.body_link_quat_w + self._get_link_quat_fn = lambda a=self._asset: a.data.body_link_quat_w.warp else: raise ValueError(f"Unsupported asset type: {self._asset.__class__.__name__}") @@ -85,6 +86,10 @@ def __init__(self, asset: BaseArticulation | BaseRigidObject | BaseRigidObjectCo self._out_force_b = wp.zeros((self.num_envs, self.num_bodies), dtype=wp.vec3f, device=self.device) self._out_torque_b = wp.zeros((self.num_envs, self.num_bodies), dtype=wp.vec3f, device=self.device) + # ProxyArray caches for the output buffers, exposed via the public properties. + self._out_force_b_ta = ProxyArray(self._out_force_b) + self._out_torque_b_ta = ProxyArray(self._out_torque_b) + # -- Index / mask helper arrays -- self._ALL_ENV_INDICES = wp.array(np.arange(self.num_envs, dtype=np.int32), dtype=wp.int32, device=self.device) self._ALL_BODY_INDICES = wp.array( @@ -164,26 +169,34 @@ def local_torque_b(self) -> wp.array: return self._local_torque_b @property - def out_force_b(self) -> wp.array: - """Composed output force [N] in the body frame, dtype ``wp.vec3f``. Shape: ``(num_envs, num_bodies)``. + def out_force_b(self) -> ProxyArray: + """Composed output force [N] in the body frame. + + Shape is ``(num_envs, num_bodies)``, dtype = ``wp.vec3f``. In torch this resolves to + ``(num_envs, num_bodies, 3)``. Use ``.warp`` for the underlying :class:`wp.array` or + ``.torch`` for a cached zero-copy :class:`torch.Tensor` view. Triggers composition from input buffers if dirty. """ self._ensure_composed() - return self._out_force_b + return self._out_force_b_ta @property - def out_torque_b(self) -> wp.array: - """Composed output torque [N·m] in the body frame, dtype ``wp.vec3f``. Shape: ``(num_envs, num_bodies)``. + def out_torque_b(self) -> ProxyArray: + """Composed output torque [N·m] in the body frame. + + Shape is ``(num_envs, num_bodies)``, dtype = ``wp.vec3f``. In torch this resolves to + ``(num_envs, num_bodies, 3)``. Use ``.warp`` for the underlying :class:`wp.array` or + ``.torch`` for a cached zero-copy :class:`torch.Tensor` view. Triggers composition from input buffers if dirty. """ self._ensure_composed() - return self._out_torque_b + return self._out_torque_b_ta @property - def composed_force(self) -> wp.array: - """Composed force at the body frame, dtype ``wp.vec3f``. Shape: ``(num_envs, num_bodies)``. + def composed_force(self) -> ProxyArray: + """Composed force in the body frame. .. deprecated:: 4.5.33 Use :attr:`out_force_b` instead. @@ -196,8 +209,8 @@ def composed_force(self) -> wp.array: return self.out_force_b @property - def composed_torque(self) -> wp.array: - """Composed torque at the body frame, dtype ``wp.vec3f``. Shape: ``(num_envs, num_bodies)``. + def composed_torque(self) -> ProxyArray: + """Composed torque in the body frame. .. deprecated:: 4.5.33 Use :attr:`out_torque_b` instead. diff --git a/source/isaaclab/test/assets/check_external_force.py b/source/isaaclab/test/assets/check_external_force.py index f038b907a1f0..f4b0da603812 100644 --- a/source/isaaclab/test/assets/check_external_force.py +++ b/source/isaaclab/test/assets/check_external_force.py @@ -96,13 +96,13 @@ def main(): sim_time = 0.0 count = 0 # reset root state - root_state = robot.data.default_root_state.clone() + root_state = robot.data.default_root_state.torch.clone() root_state[0, :2] = torch.tensor([0.0, -0.5], device=sim.device) root_state[1, :2] = torch.tensor([0.0, 0.5], device=sim.device) robot.write_root_pose_to_sim(root_state[:, :7]) robot.write_root_velocity_to_sim(root_state[:, 7:]) # reset dof state - joint_pos, joint_vel = robot.data.default_joint_pos, robot.data.default_joint_vel + joint_pos, joint_vel = robot.data.default_joint_pos.torch, robot.data.default_joint_vel.torch robot.write_joint_state_to_sim(joint_pos, joint_vel) robot.reset() # apply force @@ -112,7 +112,7 @@ def main(): # reset command print(">>>>>>>> Reset!") # apply action to the robot - robot.set_joint_position_target(robot.data.default_joint_pos.clone()) + robot.set_joint_position_target(robot.data.default_joint_pos.torch.clone()) robot.write_data_to_sim() # perform step sim.step() diff --git a/source/isaaclab/test/assets/check_fixed_base_assets.py b/source/isaaclab/test/assets/check_fixed_base_assets.py index c62c5a3334da..11f9010a7d88 100644 --- a/source/isaaclab/test/assets/check_fixed_base_assets.py +++ b/source/isaaclab/test/assets/check_fixed_base_assets.py @@ -108,13 +108,16 @@ def run_simulator(sim: sim_utils.SimulationContext, entities: dict[str, Articula # reset robots for index, robot in enumerate(entities.values()): # root state - root_state = robot.data.default_root_state.clone() + root_state = robot.data.default_root_state.torch.clone() root_state[:, :3] += origins[index] root_state[:, :2] += torch.randn_like(root_state[:, :2]) * 0.25 robot.write_root_pose_to_sim(root_state[:, :7]) robot.write_root_velocity_to_sim(root_state[:, 7:]) # joint state - joint_pos, joint_vel = robot.data.default_joint_pos.clone(), robot.data.default_joint_vel.clone() + joint_pos, joint_vel = ( + robot.data.default_joint_pos.torch.clone(), + robot.data.default_joint_vel.torch.clone(), + ) robot.write_joint_state_to_sim(joint_pos, joint_vel) # reset the internal state robot.reset() @@ -124,7 +127,7 @@ def run_simulator(sim: sim_utils.SimulationContext, entities: dict[str, Articula if count % 200 == 0: print("Name: ", name, "is_fixed_base: ", robot.is_fixed_base) # generate random joint positions - joint_pos_target = robot.data.default_joint_pos + torch.randn_like(robot.data.joint_pos) * 0.1 + joint_pos_target = robot.data.default_joint_pos.torch + torch.randn_like(robot.data.joint_pos.torch) * 0.1 # apply action to the robot robot.set_joint_position_target(joint_pos_target) # write data to sim diff --git a/source/isaaclab/test/assets/check_ridgeback_franka.py b/source/isaaclab/test/assets/check_ridgeback_franka.py index 724cd8b2a080..69ca6a7d8761 100644 --- a/source/isaaclab/test/assets/check_ridgeback_franka.py +++ b/source/isaaclab/test/assets/check_ridgeback_franka.py @@ -72,7 +72,7 @@ def add_robots() -> Articulation: def run_simulator(sim: sim_utils.SimulationContext, robot: Articulation): """Runs the simulator by applying actions to the robot at every time-step""" # dummy action - actions = robot.data.default_joint_pos.clone() + actions = robot.data.default_joint_pos.torch.clone() # Define simulation stepping sim_dt = sim.get_physics_dt() @@ -87,12 +87,15 @@ def run_simulator(sim: sim_utils.SimulationContext, robot: Articulation): sim_time = 0.0 ep_step_count = 0 # reset dof state - joint_pos, joint_vel = robot.data.default_joint_pos.clone(), robot.data.default_joint_vel.clone() + joint_pos, joint_vel = ( + robot.data.default_joint_pos.torch.clone(), + robot.data.default_joint_vel.torch.clone(), + ) robot.write_joint_state_to_sim(joint_pos, joint_vel) # reset internals robot.reset() # reset command - actions = torch.rand_like(robot.data.default_joint_pos) + robot.data.default_joint_pos + actions = torch.rand_like(robot.data.default_joint_pos.torch) + robot.data.default_joint_pos.torch # -- base actions[:, 0:3] = 0.0 # -- gripper diff --git a/source/isaaclab/test/assets/test_articulation_iface.py b/source/isaaclab/test/assets/test_articulation_iface.py index 9edc02656266..9682df92f3a5 100644 --- a/source/isaaclab/test/assets/test_articulation_iface.py +++ b/source/isaaclab/test/assets/test_articulation_iface.py @@ -465,9 +465,11 @@ def articulation_iface(request): # --------------------------------------------------------------------------- -def _check_wp_array(arr, *, expected_shape: tuple, expected_dtype: type, name: str): - """Assert that `arr` is a wp.array with the expected shape and dtype.""" - assert isinstance(arr, wp.array), f"{name}: expected wp.array, got {type(arr)}" +def _check_proxy_array(arr, *, expected_shape: tuple, expected_dtype: type, name: str): + """Assert that `arr` is a ProxyArray with the expected shape and dtype.""" + from isaaclab.utils.warp import ProxyArray + + assert isinstance(arr, ProxyArray), f"{name}: expected ProxyArray, got {type(arr)}" assert arr.shape == expected_shape, f"{name}: expected shape {expected_shape}, got {arr.shape}" assert arr.dtype == expected_dtype, f"{name}: expected dtype {expected_dtype}, got {arr.dtype}" @@ -682,7 +684,7 @@ class TestArticulationDataRootState: def test_root_link_pose_w(self, backend, num_instances, num_joints, num_bodies, device, articulation_iface): art, _ = articulation_iface art.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( art.data.root_link_pose_w, expected_shape=(num_instances,), expected_dtype=wp.transformf, @@ -695,7 +697,7 @@ def test_root_link_pose_w(self, backend, num_instances, num_joints, num_bodies, def test_root_link_vel_w(self, backend, num_instances, num_joints, num_bodies, device, articulation_iface): art, _ = articulation_iface art.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( art.data.root_link_vel_w, expected_shape=(num_instances,), expected_dtype=wp.spatial_vectorf, @@ -708,7 +710,7 @@ def test_root_link_vel_w(self, backend, num_instances, num_joints, num_bodies, d def test_root_com_pose_w(self, backend, num_instances, num_joints, num_bodies, device, articulation_iface): art, _ = articulation_iface art.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( art.data.root_com_pose_w, expected_shape=(num_instances,), expected_dtype=wp.transformf, @@ -721,7 +723,7 @@ def test_root_com_pose_w(self, backend, num_instances, num_joints, num_bodies, d def test_root_com_vel_w(self, backend, num_instances, num_joints, num_bodies, device, articulation_iface): art, _ = articulation_iface art.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( art.data.root_com_vel_w, expected_shape=(num_instances,), expected_dtype=wp.spatial_vectorf, @@ -734,7 +736,7 @@ def test_root_com_vel_w(self, backend, num_instances, num_joints, num_bodies, de def test_root_link_pos_w(self, backend, num_instances, num_joints, num_bodies, device, articulation_iface): art, _ = articulation_iface art.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( art.data.root_link_pos_w, expected_shape=(num_instances,), expected_dtype=wp.vec3f, name="root_link_pos_w" ) @@ -744,7 +746,7 @@ def test_root_link_pos_w(self, backend, num_instances, num_joints, num_bodies, d def test_root_link_quat_w(self, backend, num_instances, num_joints, num_bodies, device, articulation_iface): art, _ = articulation_iface art.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( art.data.root_link_quat_w, expected_shape=(num_instances,), expected_dtype=wp.quatf, name="root_link_quat_w" ) @@ -754,7 +756,7 @@ def test_root_link_quat_w(self, backend, num_instances, num_joints, num_bodies, def test_root_link_lin_vel_w(self, backend, num_instances, num_joints, num_bodies, device, articulation_iface): art, _ = articulation_iface art.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( art.data.root_link_lin_vel_w, expected_shape=(num_instances,), expected_dtype=wp.vec3f, @@ -767,7 +769,7 @@ def test_root_link_lin_vel_w(self, backend, num_instances, num_joints, num_bodie def test_root_link_ang_vel_w(self, backend, num_instances, num_joints, num_bodies, device, articulation_iface): art, _ = articulation_iface art.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( art.data.root_link_ang_vel_w, expected_shape=(num_instances,), expected_dtype=wp.vec3f, @@ -780,7 +782,7 @@ def test_root_link_ang_vel_w(self, backend, num_instances, num_joints, num_bodie def test_root_com_pos_w(self, backend, num_instances, num_joints, num_bodies, device, articulation_iface): art, _ = articulation_iface art.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( art.data.root_com_pos_w, expected_shape=(num_instances,), expected_dtype=wp.vec3f, name="root_com_pos_w" ) @@ -790,7 +792,7 @@ def test_root_com_pos_w(self, backend, num_instances, num_joints, num_bodies, de def test_root_com_quat_w(self, backend, num_instances, num_joints, num_bodies, device, articulation_iface): art, _ = articulation_iface art.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( art.data.root_com_quat_w, expected_shape=(num_instances,), expected_dtype=wp.quatf, name="root_com_quat_w" ) @@ -800,7 +802,7 @@ def test_root_com_quat_w(self, backend, num_instances, num_joints, num_bodies, d def test_root_com_lin_vel_w(self, backend, num_instances, num_joints, num_bodies, device, articulation_iface): art, _ = articulation_iface art.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( art.data.root_com_lin_vel_w, expected_shape=(num_instances,), expected_dtype=wp.vec3f, @@ -813,7 +815,7 @@ def test_root_com_lin_vel_w(self, backend, num_instances, num_joints, num_bodies def test_root_com_ang_vel_w(self, backend, num_instances, num_joints, num_bodies, device, articulation_iface): art, _ = articulation_iface art.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( art.data.root_com_ang_vel_w, expected_shape=(num_instances,), expected_dtype=wp.vec3f, @@ -835,7 +837,7 @@ class TestArticulationDataDerivedProperties: def test_projected_gravity_b(self, backend, num_instances, num_joints, num_bodies, device, articulation_iface): art, _ = articulation_iface art.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( art.data.projected_gravity_b, expected_shape=(num_instances,), expected_dtype=wp.vec3f, @@ -848,7 +850,7 @@ def test_projected_gravity_b(self, backend, num_instances, num_joints, num_bodie def test_heading_w(self, backend, num_instances, num_joints, num_bodies, device, articulation_iface): art, _ = articulation_iface art.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( art.data.heading_w, expected_shape=(num_instances,), expected_dtype=wp.float32, name="heading_w" ) @@ -858,7 +860,7 @@ def test_heading_w(self, backend, num_instances, num_joints, num_bodies, device, def test_root_link_lin_vel_b(self, backend, num_instances, num_joints, num_bodies, device, articulation_iface): art, _ = articulation_iface art.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( art.data.root_link_lin_vel_b, expected_shape=(num_instances,), expected_dtype=wp.vec3f, @@ -871,7 +873,7 @@ def test_root_link_lin_vel_b(self, backend, num_instances, num_joints, num_bodie def test_root_link_ang_vel_b(self, backend, num_instances, num_joints, num_bodies, device, articulation_iface): art, _ = articulation_iface art.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( art.data.root_link_ang_vel_b, expected_shape=(num_instances,), expected_dtype=wp.vec3f, @@ -884,7 +886,7 @@ def test_root_link_ang_vel_b(self, backend, num_instances, num_joints, num_bodie def test_root_com_lin_vel_b(self, backend, num_instances, num_joints, num_bodies, device, articulation_iface): art, _ = articulation_iface art.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( art.data.root_com_lin_vel_b, expected_shape=(num_instances,), expected_dtype=wp.vec3f, @@ -897,7 +899,7 @@ def test_root_com_lin_vel_b(self, backend, num_instances, num_joints, num_bodies def test_root_com_ang_vel_b(self, backend, num_instances, num_joints, num_bodies, device, articulation_iface): art, _ = articulation_iface art.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( art.data.root_com_ang_vel_b, expected_shape=(num_instances,), expected_dtype=wp.vec3f, @@ -919,7 +921,7 @@ class TestArticulationDataBodyState: def test_body_link_pose_w(self, backend, num_instances, num_joints, num_bodies, device, articulation_iface): art, _ = articulation_iface art.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( art.data.body_link_pose_w, expected_shape=(num_instances, num_bodies), expected_dtype=wp.transformf, @@ -932,7 +934,7 @@ def test_body_link_pose_w(self, backend, num_instances, num_joints, num_bodies, def test_body_link_vel_w(self, backend, num_instances, num_joints, num_bodies, device, articulation_iface): art, _ = articulation_iface art.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( art.data.body_link_vel_w, expected_shape=(num_instances, num_bodies), expected_dtype=wp.spatial_vectorf, @@ -945,7 +947,7 @@ def test_body_link_vel_w(self, backend, num_instances, num_joints, num_bodies, d def test_body_com_pose_w(self, backend, num_instances, num_joints, num_bodies, device, articulation_iface): art, _ = articulation_iface art.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( art.data.body_com_pose_w, expected_shape=(num_instances, num_bodies), expected_dtype=wp.transformf, @@ -958,7 +960,7 @@ def test_body_com_pose_w(self, backend, num_instances, num_joints, num_bodies, d def test_body_com_vel_w(self, backend, num_instances, num_joints, num_bodies, device, articulation_iface): art, _ = articulation_iface art.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( art.data.body_com_vel_w, expected_shape=(num_instances, num_bodies), expected_dtype=wp.spatial_vectorf, @@ -971,7 +973,7 @@ def test_body_com_vel_w(self, backend, num_instances, num_joints, num_bodies, de def test_body_com_acc_w(self, backend, num_instances, num_joints, num_bodies, device, articulation_iface): art, _ = articulation_iface art.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( art.data.body_com_acc_w, expected_shape=(num_instances, num_bodies), expected_dtype=wp.spatial_vectorf, @@ -986,7 +988,7 @@ def test_body_com_pose_b(self, backend, num_instances, num_joints, num_bodies, d pytest.xfail("Newton only stores CoM position, not orientation") art, _ = articulation_iface art.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( art.data.body_com_pose_b, expected_shape=(num_instances, num_bodies), expected_dtype=wp.transformf, @@ -999,7 +1001,7 @@ def test_body_com_pose_b(self, backend, num_instances, num_joints, num_bodies, d def test_body_mass(self, backend, num_instances, num_joints, num_bodies, device, articulation_iface): art, _ = articulation_iface art.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( art.data.body_mass, expected_shape=(num_instances, num_bodies), expected_dtype=wp.float32, name="body_mass" ) @@ -1009,7 +1011,7 @@ def test_body_mass(self, backend, num_instances, num_joints, num_bodies, device, def test_body_inertia(self, backend, num_instances, num_joints, num_bodies, device, articulation_iface): art, _ = articulation_iface art.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( art.data.body_inertia, expected_shape=(num_instances, num_bodies, 9), expected_dtype=wp.float32, @@ -1026,7 +1028,7 @@ def test_body_incoming_joint_wrench_b( pytest.xfail("Newton does not support joint wrench reporting") art, _ = articulation_iface art.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( art.data.body_incoming_joint_wrench_b, expected_shape=(num_instances, num_bodies), expected_dtype=wp.spatial_vectorf, @@ -1039,7 +1041,7 @@ def test_body_incoming_joint_wrench_b( def test_body_link_pos_w(self, backend, num_instances, num_joints, num_bodies, device, articulation_iface): art, _ = articulation_iface art.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( art.data.body_link_pos_w, expected_shape=(num_instances, num_bodies), expected_dtype=wp.vec3f, @@ -1052,7 +1054,7 @@ def test_body_link_pos_w(self, backend, num_instances, num_joints, num_bodies, d def test_body_link_quat_w(self, backend, num_instances, num_joints, num_bodies, device, articulation_iface): art, _ = articulation_iface art.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( art.data.body_link_quat_w, expected_shape=(num_instances, num_bodies), expected_dtype=wp.quatf, @@ -1065,7 +1067,7 @@ def test_body_link_quat_w(self, backend, num_instances, num_joints, num_bodies, def test_body_link_lin_vel_w(self, backend, num_instances, num_joints, num_bodies, device, articulation_iface): art, _ = articulation_iface art.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( art.data.body_link_lin_vel_w, expected_shape=(num_instances, num_bodies), expected_dtype=wp.vec3f, @@ -1078,7 +1080,7 @@ def test_body_link_lin_vel_w(self, backend, num_instances, num_joints, num_bodie def test_body_link_ang_vel_w(self, backend, num_instances, num_joints, num_bodies, device, articulation_iface): art, _ = articulation_iface art.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( art.data.body_link_ang_vel_w, expected_shape=(num_instances, num_bodies), expected_dtype=wp.vec3f, @@ -1091,7 +1093,7 @@ def test_body_link_ang_vel_w(self, backend, num_instances, num_joints, num_bodie def test_body_com_pos_w(self, backend, num_instances, num_joints, num_bodies, device, articulation_iface): art, _ = articulation_iface art.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( art.data.body_com_pos_w, expected_shape=(num_instances, num_bodies), expected_dtype=wp.vec3f, @@ -1104,7 +1106,7 @@ def test_body_com_pos_w(self, backend, num_instances, num_joints, num_bodies, de def test_body_com_quat_w(self, backend, num_instances, num_joints, num_bodies, device, articulation_iface): art, _ = articulation_iface art.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( art.data.body_com_quat_w, expected_shape=(num_instances, num_bodies), expected_dtype=wp.quatf, @@ -1117,7 +1119,7 @@ def test_body_com_quat_w(self, backend, num_instances, num_joints, num_bodies, d def test_body_com_pos_b(self, backend, num_instances, num_joints, num_bodies, device, articulation_iface): art, _ = articulation_iface art.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( art.data.body_com_pos_b, expected_shape=(num_instances, num_bodies), expected_dtype=wp.vec3f, @@ -1132,7 +1134,7 @@ def test_body_com_quat_b(self, backend, num_instances, num_joints, num_bodies, d pytest.xfail("Newton only stores CoM position, not orientation") art, _ = articulation_iface art.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( art.data.body_com_quat_b, expected_shape=(num_instances, num_bodies), expected_dtype=wp.quatf, @@ -1154,7 +1156,7 @@ class TestArticulationDataJointState: def test_joint_pos(self, backend, num_instances, num_joints, num_bodies, device, articulation_iface): art, _ = articulation_iface art.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( art.data.joint_pos, expected_shape=(num_instances, num_joints), expected_dtype=wp.float32, name="joint_pos" ) @@ -1164,7 +1166,7 @@ def test_joint_pos(self, backend, num_instances, num_joints, num_bodies, device, def test_joint_vel(self, backend, num_instances, num_joints, num_bodies, device, articulation_iface): art, _ = articulation_iface art.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( art.data.joint_vel, expected_shape=(num_instances, num_joints), expected_dtype=wp.float32, name="joint_vel" ) @@ -1174,7 +1176,7 @@ def test_joint_vel(self, backend, num_instances, num_joints, num_bodies, device, def test_joint_acc(self, backend, num_instances, num_joints, num_bodies, device, articulation_iface): art, _ = articulation_iface art.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( art.data.joint_acc, expected_shape=(num_instances, num_joints), expected_dtype=wp.float32, name="joint_acc" ) @@ -1184,7 +1186,7 @@ def test_joint_acc(self, backend, num_instances, num_joints, num_bodies, device, def test_joint_stiffness(self, backend, num_instances, num_joints, num_bodies, device, articulation_iface): art, _ = articulation_iface art.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( art.data.joint_stiffness, expected_shape=(num_instances, num_joints), expected_dtype=wp.float32, @@ -1197,7 +1199,7 @@ def test_joint_stiffness(self, backend, num_instances, num_joints, num_bodies, d def test_joint_damping(self, backend, num_instances, num_joints, num_bodies, device, articulation_iface): art, _ = articulation_iface art.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( art.data.joint_damping, expected_shape=(num_instances, num_joints), expected_dtype=wp.float32, @@ -1210,7 +1212,7 @@ def test_joint_damping(self, backend, num_instances, num_joints, num_bodies, dev def test_joint_armature(self, backend, num_instances, num_joints, num_bodies, device, articulation_iface): art, _ = articulation_iface art.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( art.data.joint_armature, expected_shape=(num_instances, num_joints), expected_dtype=wp.float32, @@ -1223,7 +1225,7 @@ def test_joint_armature(self, backend, num_instances, num_joints, num_bodies, de def test_joint_friction_coeff(self, backend, num_instances, num_joints, num_bodies, device, articulation_iface): art, _ = articulation_iface art.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( art.data.joint_friction_coeff, expected_shape=(num_instances, num_joints), expected_dtype=wp.float32, @@ -1236,7 +1238,7 @@ def test_joint_friction_coeff(self, backend, num_instances, num_joints, num_bodi def test_joint_pos_limits(self, backend, num_instances, num_joints, num_bodies, device, articulation_iface): art, _ = articulation_iface art.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( art.data.joint_pos_limits, expected_shape=(num_instances, num_joints), expected_dtype=wp.vec2f, @@ -1249,7 +1251,7 @@ def test_joint_pos_limits(self, backend, num_instances, num_joints, num_bodies, def test_joint_vel_limits(self, backend, num_instances, num_joints, num_bodies, device, articulation_iface): art, _ = articulation_iface art.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( art.data.joint_vel_limits, expected_shape=(num_instances, num_joints), expected_dtype=wp.float32, @@ -1262,7 +1264,7 @@ def test_joint_vel_limits(self, backend, num_instances, num_joints, num_bodies, def test_joint_effort_limits(self, backend, num_instances, num_joints, num_bodies, device, articulation_iface): art, _ = articulation_iface art.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( art.data.joint_effort_limits, expected_shape=(num_instances, num_joints), expected_dtype=wp.float32, @@ -1275,7 +1277,7 @@ def test_joint_effort_limits(self, backend, num_instances, num_joints, num_bodie def test_soft_joint_pos_limits(self, backend, num_instances, num_joints, num_bodies, device, articulation_iface): art, _ = articulation_iface art.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( art.data.soft_joint_pos_limits, expected_shape=(num_instances, num_joints), expected_dtype=wp.vec2f, @@ -1297,7 +1299,7 @@ class TestArticulationDataDefaults: def test_default_root_pose(self, backend, num_instances, num_joints, num_bodies, device, articulation_iface): art, _ = articulation_iface art.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( art.data.default_root_pose, expected_shape=(num_instances,), expected_dtype=wp.transformf, @@ -1310,7 +1312,7 @@ def test_default_root_pose(self, backend, num_instances, num_joints, num_bodies, def test_default_root_vel(self, backend, num_instances, num_joints, num_bodies, device, articulation_iface): art, _ = articulation_iface art.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( art.data.default_root_vel, expected_shape=(num_instances,), expected_dtype=wp.spatial_vectorf, @@ -1323,7 +1325,7 @@ def test_default_root_vel(self, backend, num_instances, num_joints, num_bodies, def test_default_joint_pos(self, backend, num_instances, num_joints, num_bodies, device, articulation_iface): art, _ = articulation_iface art.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( art.data.default_joint_pos, expected_shape=(num_instances, num_joints), expected_dtype=wp.float32, @@ -1336,7 +1338,7 @@ def test_default_joint_pos(self, backend, num_instances, num_joints, num_bodies, def test_default_joint_vel(self, backend, num_instances, num_joints, num_bodies, device, articulation_iface): art, _ = articulation_iface art.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( art.data.default_joint_vel, expected_shape=(num_instances, num_joints), expected_dtype=wp.float32, @@ -1349,7 +1351,7 @@ def test_default_joint_vel(self, backend, num_instances, num_joints, num_bodies, def test_joint_pos_target(self, backend, num_instances, num_joints, num_bodies, device, articulation_iface): art, _ = articulation_iface art.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( art.data.joint_pos_target, expected_shape=(num_instances, num_joints), expected_dtype=wp.float32, @@ -1362,7 +1364,7 @@ def test_joint_pos_target(self, backend, num_instances, num_joints, num_bodies, def test_joint_vel_target(self, backend, num_instances, num_joints, num_bodies, device, articulation_iface): art, _ = articulation_iface art.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( art.data.joint_vel_target, expected_shape=(num_instances, num_joints), expected_dtype=wp.float32, @@ -1375,7 +1377,7 @@ def test_joint_vel_target(self, backend, num_instances, num_joints, num_bodies, def test_joint_effort_target(self, backend, num_instances, num_joints, num_bodies, device, articulation_iface): art, _ = articulation_iface art.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( art.data.joint_effort_target, expected_shape=(num_instances, num_joints), expected_dtype=wp.float32, @@ -2186,7 +2188,7 @@ def test_fixed_tendon_stiffness( ): art, _ = articulation_iface art.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( art.data.fixed_tendon_stiffness, expected_shape=(num_instances, num_fixed_tendons), expected_dtype=wp.float32, @@ -2209,7 +2211,7 @@ def test_fixed_tendon_damping( ): art, _ = articulation_iface art.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( art.data.fixed_tendon_damping, expected_shape=(num_instances, num_fixed_tendons), expected_dtype=wp.float32, @@ -2232,7 +2234,7 @@ def test_fixed_tendon_limit_stiffness( ): art, _ = articulation_iface art.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( art.data.fixed_tendon_limit_stiffness, expected_shape=(num_instances, num_fixed_tendons), expected_dtype=wp.float32, @@ -2255,7 +2257,7 @@ def test_fixed_tendon_rest_length( ): art, _ = articulation_iface art.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( art.data.fixed_tendon_rest_length, expected_shape=(num_instances, num_fixed_tendons), expected_dtype=wp.float32, @@ -2278,7 +2280,7 @@ def test_fixed_tendon_offset( ): art, _ = articulation_iface art.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( art.data.fixed_tendon_offset, expected_shape=(num_instances, num_fixed_tendons), expected_dtype=wp.float32, @@ -2301,8 +2303,10 @@ def test_fixed_tendon_pos_limits( ): art, _ = articulation_iface art.data.update(dt=0.01) + from isaaclab.utils.warp import ProxyArray + arr = art.data.fixed_tendon_pos_limits - assert isinstance(arr, wp.array), f"fixed_tendon_pos_limits: expected wp.array, got {type(arr)}" + assert isinstance(arr, ProxyArray), f"fixed_tendon_pos_limits: expected ProxyArray, got {type(arr)}" if num_fixed_tendons == 0: # When no tendons, shape is (N, 0, 2) float32 assert arr.shape == (num_instances, 0, 2) @@ -2332,7 +2336,7 @@ def test_spatial_tendon_stiffness( if num_spatial_tendons == 0: pytest.skip("No spatial tendons configured") art.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( art.data.spatial_tendon_stiffness, expected_shape=(num_instances, num_spatial_tendons), expected_dtype=wp.float32, @@ -2357,7 +2361,7 @@ def test_spatial_tendon_damping( if num_spatial_tendons == 0: pytest.skip("No spatial tendons configured") art.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( art.data.spatial_tendon_damping, expected_shape=(num_instances, num_spatial_tendons), expected_dtype=wp.float32, @@ -2382,7 +2386,7 @@ def test_spatial_tendon_limit_stiffness( if num_spatial_tendons == 0: pytest.skip("No spatial tendons configured") art.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( art.data.spatial_tendon_limit_stiffness, expected_shape=(num_instances, num_spatial_tendons), expected_dtype=wp.float32, @@ -2407,7 +2411,7 @@ def test_spatial_tendon_offset( if num_spatial_tendons == 0: pytest.skip("No spatial tendons configured") art.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( art.data.spatial_tendon_offset, expected_shape=(num_instances, num_spatial_tendons), expected_dtype=wp.float32, diff --git a/source/isaaclab/test/assets/test_rigid_object_collection_iface.py b/source/isaaclab/test/assets/test_rigid_object_collection_iface.py index c047329d0885..4d067f44c1ce 100644 --- a/source/isaaclab/test/assets/test_rigid_object_collection_iface.py +++ b/source/isaaclab/test/assets/test_rigid_object_collection_iface.py @@ -254,9 +254,11 @@ def collection_iface(request): # --------------------------------------------------------------------------- -def _check_wp_array(arr, *, expected_shape: tuple, expected_dtype: type, name: str): - """Assert that `arr` is a wp.array with the expected shape and dtype.""" - assert isinstance(arr, wp.array), f"{name}: expected wp.array, got {type(arr)}" +def _check_proxy_array(arr, *, expected_shape: tuple, expected_dtype: type, name: str): + """Assert that `arr` is a ProxyArray with the expected shape and dtype.""" + from isaaclab.utils.warp import ProxyArray + + assert isinstance(arr, ProxyArray), f"{name}: expected ProxyArray, got {type(arr)}" assert arr.shape == expected_shape, f"{name}: expected shape {expected_shape}, got {arr.shape}" assert arr.dtype == expected_dtype, f"{name}: expected dtype {expected_dtype}, got {arr.dtype}" @@ -443,7 +445,7 @@ class TestCollectionDataBodyState: def test_body_link_pose_w(self, backend, num_instances, num_bodies, device, collection_iface): obj, _ = collection_iface obj.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( obj.data.body_link_pose_w, expected_shape=(num_instances, num_bodies), expected_dtype=wp.transformf, @@ -457,7 +459,7 @@ def test_body_link_pose_w(self, backend, num_instances, num_bodies, device, coll def test_body_link_vel_w(self, backend, num_instances, num_bodies, device, collection_iface): obj, _ = collection_iface obj.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( obj.data.body_link_vel_w, expected_shape=(num_instances, num_bodies), expected_dtype=wp.spatial_vectorf, @@ -471,7 +473,7 @@ def test_body_link_vel_w(self, backend, num_instances, num_bodies, device, colle def test_body_com_pose_w(self, backend, num_instances, num_bodies, device, collection_iface): obj, _ = collection_iface obj.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( obj.data.body_com_pose_w, expected_shape=(num_instances, num_bodies), expected_dtype=wp.transformf, @@ -485,7 +487,7 @@ def test_body_com_pose_w(self, backend, num_instances, num_bodies, device, colle def test_body_com_vel_w(self, backend, num_instances, num_bodies, device, collection_iface): obj, _ = collection_iface obj.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( obj.data.body_com_vel_w, expected_shape=(num_instances, num_bodies), expected_dtype=wp.spatial_vectorf, @@ -499,7 +501,7 @@ def test_body_com_vel_w(self, backend, num_instances, num_bodies, device, collec def test_body_com_acc_w(self, backend, num_instances, num_bodies, device, collection_iface): obj, _ = collection_iface obj.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( obj.data.body_com_acc_w, expected_shape=(num_instances, num_bodies), expected_dtype=wp.spatial_vectorf, @@ -513,7 +515,7 @@ def test_body_com_acc_w(self, backend, num_instances, num_bodies, device, collec def test_body_com_pose_b(self, backend, num_instances, num_bodies, device, collection_iface): obj, _ = collection_iface obj.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( obj.data.body_com_pose_b, expected_shape=(num_instances, num_bodies), expected_dtype=wp.transformf, @@ -536,7 +538,7 @@ class TestCollectionDataSliced: def test_body_link_pos_w(self, backend, num_instances, num_bodies, device, collection_iface): obj, _ = collection_iface obj.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( obj.data.body_link_pos_w, expected_shape=(num_instances, num_bodies), expected_dtype=wp.vec3f, @@ -550,7 +552,7 @@ def test_body_link_pos_w(self, backend, num_instances, num_bodies, device, colle def test_body_link_quat_w(self, backend, num_instances, num_bodies, device, collection_iface): obj, _ = collection_iface obj.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( obj.data.body_link_quat_w, expected_shape=(num_instances, num_bodies), expected_dtype=wp.quatf, @@ -564,7 +566,7 @@ def test_body_link_quat_w(self, backend, num_instances, num_bodies, device, coll def test_body_link_lin_vel_w(self, backend, num_instances, num_bodies, device, collection_iface): obj, _ = collection_iface obj.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( obj.data.body_link_lin_vel_w, expected_shape=(num_instances, num_bodies), expected_dtype=wp.vec3f, @@ -578,7 +580,7 @@ def test_body_link_lin_vel_w(self, backend, num_instances, num_bodies, device, c def test_body_link_ang_vel_w(self, backend, num_instances, num_bodies, device, collection_iface): obj, _ = collection_iface obj.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( obj.data.body_link_ang_vel_w, expected_shape=(num_instances, num_bodies), expected_dtype=wp.vec3f, @@ -592,7 +594,7 @@ def test_body_link_ang_vel_w(self, backend, num_instances, num_bodies, device, c def test_body_com_pos_w(self, backend, num_instances, num_bodies, device, collection_iface): obj, _ = collection_iface obj.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( obj.data.body_com_pos_w, expected_shape=(num_instances, num_bodies), expected_dtype=wp.vec3f, @@ -606,7 +608,7 @@ def test_body_com_pos_w(self, backend, num_instances, num_bodies, device, collec def test_body_com_quat_w(self, backend, num_instances, num_bodies, device, collection_iface): obj, _ = collection_iface obj.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( obj.data.body_com_quat_w, expected_shape=(num_instances, num_bodies), expected_dtype=wp.quatf, @@ -620,7 +622,7 @@ def test_body_com_quat_w(self, backend, num_instances, num_bodies, device, colle def test_body_com_lin_vel_w(self, backend, num_instances, num_bodies, device, collection_iface): obj, _ = collection_iface obj.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( obj.data.body_com_lin_vel_w, expected_shape=(num_instances, num_bodies), expected_dtype=wp.vec3f, @@ -634,7 +636,7 @@ def test_body_com_lin_vel_w(self, backend, num_instances, num_bodies, device, co def test_body_com_ang_vel_w(self, backend, num_instances, num_bodies, device, collection_iface): obj, _ = collection_iface obj.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( obj.data.body_com_ang_vel_w, expected_shape=(num_instances, num_bodies), expected_dtype=wp.vec3f, @@ -648,7 +650,7 @@ def test_body_com_ang_vel_w(self, backend, num_instances, num_bodies, device, co def test_body_com_lin_acc_w(self, backend, num_instances, num_bodies, device, collection_iface): obj, _ = collection_iface obj.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( obj.data.body_com_lin_acc_w, expected_shape=(num_instances, num_bodies), expected_dtype=wp.vec3f, @@ -662,7 +664,7 @@ def test_body_com_lin_acc_w(self, backend, num_instances, num_bodies, device, co def test_body_com_ang_acc_w(self, backend, num_instances, num_bodies, device, collection_iface): obj, _ = collection_iface obj.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( obj.data.body_com_ang_acc_w, expected_shape=(num_instances, num_bodies), expected_dtype=wp.vec3f, @@ -676,7 +678,7 @@ def test_body_com_ang_acc_w(self, backend, num_instances, num_bodies, device, co def test_body_com_pos_b(self, backend, num_instances, num_bodies, device, collection_iface): obj, _ = collection_iface obj.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( obj.data.body_com_pos_b, expected_shape=(num_instances, num_bodies), expected_dtype=wp.vec3f, @@ -690,7 +692,7 @@ def test_body_com_pos_b(self, backend, num_instances, num_bodies, device, collec def test_body_com_quat_b(self, backend, num_instances, num_bodies, device, collection_iface): obj, _ = collection_iface obj.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( obj.data.body_com_quat_b, expected_shape=(num_instances, num_bodies), expected_dtype=wp.quatf, @@ -713,7 +715,7 @@ class TestCollectionDataDerived: def test_projected_gravity_b(self, backend, num_instances, num_bodies, device, collection_iface): obj, _ = collection_iface obj.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( obj.data.projected_gravity_b, expected_shape=(num_instances, num_bodies), expected_dtype=wp.vec3f, @@ -727,7 +729,7 @@ def test_projected_gravity_b(self, backend, num_instances, num_bodies, device, c def test_heading_w(self, backend, num_instances, num_bodies, device, collection_iface): obj, _ = collection_iface obj.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( obj.data.heading_w, expected_shape=(num_instances, num_bodies), expected_dtype=wp.float32, name="heading_w" ) @@ -738,7 +740,7 @@ def test_heading_w(self, backend, num_instances, num_bodies, device, collection_ def test_body_link_lin_vel_b(self, backend, num_instances, num_bodies, device, collection_iface): obj, _ = collection_iface obj.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( obj.data.body_link_lin_vel_b, expected_shape=(num_instances, num_bodies), expected_dtype=wp.vec3f, @@ -752,7 +754,7 @@ def test_body_link_lin_vel_b(self, backend, num_instances, num_bodies, device, c def test_body_link_ang_vel_b(self, backend, num_instances, num_bodies, device, collection_iface): obj, _ = collection_iface obj.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( obj.data.body_link_ang_vel_b, expected_shape=(num_instances, num_bodies), expected_dtype=wp.vec3f, @@ -766,7 +768,7 @@ def test_body_link_ang_vel_b(self, backend, num_instances, num_bodies, device, c def test_body_com_lin_vel_b(self, backend, num_instances, num_bodies, device, collection_iface): obj, _ = collection_iface obj.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( obj.data.body_com_lin_vel_b, expected_shape=(num_instances, num_bodies), expected_dtype=wp.vec3f, @@ -780,7 +782,7 @@ def test_body_com_lin_vel_b(self, backend, num_instances, num_bodies, device, co def test_body_com_ang_vel_b(self, backend, num_instances, num_bodies, device, collection_iface): obj, _ = collection_iface obj.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( obj.data.body_com_ang_vel_b, expected_shape=(num_instances, num_bodies), expected_dtype=wp.vec3f, @@ -803,7 +805,7 @@ class TestCollectionDataMass: def test_body_mass(self, backend, num_instances, num_bodies, device, collection_iface): obj, _ = collection_iface obj.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( obj.data.body_mass, expected_shape=(num_instances, num_bodies), expected_dtype=wp.float32, name="body_mass" ) @@ -814,7 +816,7 @@ def test_body_mass(self, backend, num_instances, num_bodies, device, collection_ def test_body_inertia(self, backend, num_instances, num_bodies, device, collection_iface): obj, _ = collection_iface obj.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( obj.data.body_inertia, expected_shape=(num_instances, num_bodies, 9), expected_dtype=wp.float32, @@ -837,7 +839,7 @@ class TestCollectionDataDefaults: def test_default_body_pose(self, backend, num_instances, num_bodies, device, collection_iface): obj, _ = collection_iface obj.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( obj.data.default_body_pose, expected_shape=(num_instances, num_bodies), expected_dtype=wp.transformf, @@ -851,7 +853,7 @@ def test_default_body_pose(self, backend, num_instances, num_bodies, device, col def test_default_body_vel(self, backend, num_instances, num_bodies, device, collection_iface): obj, _ = collection_iface obj.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( obj.data.default_body_vel, expected_shape=(num_instances, num_bodies), expected_dtype=wp.spatial_vectorf, diff --git a/source/isaaclab/test/assets/test_rigid_object_iface.py b/source/isaaclab/test/assets/test_rigid_object_iface.py index 085005627b79..c7e01ad8ada7 100644 --- a/source/isaaclab/test/assets/test_rigid_object_iface.py +++ b/source/isaaclab/test/assets/test_rigid_object_iface.py @@ -231,9 +231,11 @@ def rigid_object_iface(request): # --------------------------------------------------------------------------- -def _check_wp_array(arr, *, expected_shape: tuple, expected_dtype: type, name: str): - """Assert that `arr` is a wp.array with the expected shape and dtype.""" - assert isinstance(arr, wp.array), f"{name}: expected wp.array, got {type(arr)}" +def _check_proxy_array(arr, *, expected_shape: tuple, expected_dtype: type, name: str): + """Assert that `arr` is a ProxyArray with the expected shape and dtype.""" + from isaaclab.utils.warp import ProxyArray + + assert isinstance(arr, ProxyArray), f"{name}: expected ProxyArray, got {type(arr)}" assert arr.shape == expected_shape, f"{name}: expected shape {expected_shape}, got {arr.shape}" assert arr.dtype == expected_dtype, f"{name}: expected dtype {expected_dtype}, got {arr.dtype}" @@ -333,7 +335,7 @@ class TestRigidObjectDataRootState: def test_root_link_pose_w(self, backend, num_instances, device, rigid_object_iface): obj, _ = rigid_object_iface obj.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( obj.data.root_link_pose_w, expected_shape=(num_instances,), expected_dtype=wp.transformf, @@ -346,7 +348,7 @@ def test_root_link_pose_w(self, backend, num_instances, device, rigid_object_ifa def test_root_link_vel_w(self, backend, num_instances, device, rigid_object_iface): obj, _ = rigid_object_iface obj.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( obj.data.root_link_vel_w, expected_shape=(num_instances,), expected_dtype=wp.spatial_vectorf, @@ -359,7 +361,7 @@ def test_root_link_vel_w(self, backend, num_instances, device, rigid_object_ifac def test_root_com_pose_w(self, backend, num_instances, device, rigid_object_iface): obj, _ = rigid_object_iface obj.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( obj.data.root_com_pose_w, expected_shape=(num_instances,), expected_dtype=wp.transformf, @@ -372,7 +374,7 @@ def test_root_com_pose_w(self, backend, num_instances, device, rigid_object_ifac def test_root_com_vel_w(self, backend, num_instances, device, rigid_object_iface): obj, _ = rigid_object_iface obj.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( obj.data.root_com_vel_w, expected_shape=(num_instances,), expected_dtype=wp.spatial_vectorf, @@ -385,7 +387,7 @@ def test_root_com_vel_w(self, backend, num_instances, device, rigid_object_iface def test_root_link_pos_w(self, backend, num_instances, device, rigid_object_iface): obj, _ = rigid_object_iface obj.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( obj.data.root_link_pos_w, expected_shape=(num_instances,), expected_dtype=wp.vec3f, name="root_link_pos_w" ) @@ -395,7 +397,7 @@ def test_root_link_pos_w(self, backend, num_instances, device, rigid_object_ifac def test_root_link_quat_w(self, backend, num_instances, device, rigid_object_iface): obj, _ = rigid_object_iface obj.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( obj.data.root_link_quat_w, expected_shape=(num_instances,), expected_dtype=wp.quatf, name="root_link_quat_w" ) @@ -405,7 +407,7 @@ def test_root_link_quat_w(self, backend, num_instances, device, rigid_object_ifa def test_root_link_lin_vel_w(self, backend, num_instances, device, rigid_object_iface): obj, _ = rigid_object_iface obj.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( obj.data.root_link_lin_vel_w, expected_shape=(num_instances,), expected_dtype=wp.vec3f, @@ -418,7 +420,7 @@ def test_root_link_lin_vel_w(self, backend, num_instances, device, rigid_object_ def test_root_link_ang_vel_w(self, backend, num_instances, device, rigid_object_iface): obj, _ = rigid_object_iface obj.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( obj.data.root_link_ang_vel_w, expected_shape=(num_instances,), expected_dtype=wp.vec3f, @@ -431,7 +433,7 @@ def test_root_link_ang_vel_w(self, backend, num_instances, device, rigid_object_ def test_root_com_pos_w(self, backend, num_instances, device, rigid_object_iface): obj, _ = rigid_object_iface obj.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( obj.data.root_com_pos_w, expected_shape=(num_instances,), expected_dtype=wp.vec3f, name="root_com_pos_w" ) @@ -441,7 +443,7 @@ def test_root_com_pos_w(self, backend, num_instances, device, rigid_object_iface def test_root_com_quat_w(self, backend, num_instances, device, rigid_object_iface): obj, _ = rigid_object_iface obj.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( obj.data.root_com_quat_w, expected_shape=(num_instances,), expected_dtype=wp.quatf, name="root_com_quat_w" ) @@ -451,7 +453,7 @@ def test_root_com_quat_w(self, backend, num_instances, device, rigid_object_ifac def test_root_com_lin_vel_w(self, backend, num_instances, device, rigid_object_iface): obj, _ = rigid_object_iface obj.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( obj.data.root_com_lin_vel_w, expected_shape=(num_instances,), expected_dtype=wp.vec3f, @@ -464,7 +466,7 @@ def test_root_com_lin_vel_w(self, backend, num_instances, device, rigid_object_i def test_root_com_ang_vel_w(self, backend, num_instances, device, rigid_object_iface): obj, _ = rigid_object_iface obj.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( obj.data.root_com_ang_vel_w, expected_shape=(num_instances,), expected_dtype=wp.vec3f, @@ -486,7 +488,7 @@ class TestRigidObjectDataDerivedProperties: def test_projected_gravity_b(self, backend, num_instances, device, rigid_object_iface): obj, _ = rigid_object_iface obj.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( obj.data.projected_gravity_b, expected_shape=(num_instances,), expected_dtype=wp.vec3f, @@ -499,7 +501,7 @@ def test_projected_gravity_b(self, backend, num_instances, device, rigid_object_ def test_heading_w(self, backend, num_instances, device, rigid_object_iface): obj, _ = rigid_object_iface obj.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( obj.data.heading_w, expected_shape=(num_instances,), expected_dtype=wp.float32, name="heading_w" ) @@ -509,7 +511,7 @@ def test_heading_w(self, backend, num_instances, device, rigid_object_iface): def test_root_link_lin_vel_b(self, backend, num_instances, device, rigid_object_iface): obj, _ = rigid_object_iface obj.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( obj.data.root_link_lin_vel_b, expected_shape=(num_instances,), expected_dtype=wp.vec3f, @@ -522,7 +524,7 @@ def test_root_link_lin_vel_b(self, backend, num_instances, device, rigid_object_ def test_root_link_ang_vel_b(self, backend, num_instances, device, rigid_object_iface): obj, _ = rigid_object_iface obj.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( obj.data.root_link_ang_vel_b, expected_shape=(num_instances,), expected_dtype=wp.vec3f, @@ -535,7 +537,7 @@ def test_root_link_ang_vel_b(self, backend, num_instances, device, rigid_object_ def test_root_com_lin_vel_b(self, backend, num_instances, device, rigid_object_iface): obj, _ = rigid_object_iface obj.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( obj.data.root_com_lin_vel_b, expected_shape=(num_instances,), expected_dtype=wp.vec3f, @@ -548,7 +550,7 @@ def test_root_com_lin_vel_b(self, backend, num_instances, device, rigid_object_i def test_root_com_ang_vel_b(self, backend, num_instances, device, rigid_object_iface): obj, _ = rigid_object_iface obj.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( obj.data.root_com_ang_vel_b, expected_shape=(num_instances,), expected_dtype=wp.vec3f, @@ -570,7 +572,7 @@ class TestRigidObjectDataBodyState: def test_body_link_pose_w(self, backend, num_instances, device, rigid_object_iface): obj, _ = rigid_object_iface obj.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( obj.data.body_link_pose_w, expected_shape=(num_instances, 1), expected_dtype=wp.transformf, @@ -583,7 +585,7 @@ def test_body_link_pose_w(self, backend, num_instances, device, rigid_object_ifa def test_body_link_vel_w(self, backend, num_instances, device, rigid_object_iface): obj, _ = rigid_object_iface obj.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( obj.data.body_link_vel_w, expected_shape=(num_instances, 1), expected_dtype=wp.spatial_vectorf, @@ -596,7 +598,7 @@ def test_body_link_vel_w(self, backend, num_instances, device, rigid_object_ifac def test_body_com_pose_w(self, backend, num_instances, device, rigid_object_iface): obj, _ = rigid_object_iface obj.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( obj.data.body_com_pose_w, expected_shape=(num_instances, 1), expected_dtype=wp.transformf, @@ -609,7 +611,7 @@ def test_body_com_pose_w(self, backend, num_instances, device, rigid_object_ifac def test_body_com_vel_w(self, backend, num_instances, device, rigid_object_iface): obj, _ = rigid_object_iface obj.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( obj.data.body_com_vel_w, expected_shape=(num_instances, 1), expected_dtype=wp.spatial_vectorf, @@ -622,7 +624,7 @@ def test_body_com_vel_w(self, backend, num_instances, device, rigid_object_iface def test_body_com_acc_w(self, backend, num_instances, device, rigid_object_iface): obj, _ = rigid_object_iface obj.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( obj.data.body_com_acc_w, expected_shape=(num_instances, 1), expected_dtype=wp.spatial_vectorf, @@ -635,7 +637,7 @@ def test_body_com_acc_w(self, backend, num_instances, device, rigid_object_iface def test_body_com_pose_b(self, backend, num_instances, device, rigid_object_iface): obj, _ = rigid_object_iface obj.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( obj.data.body_com_pose_b, expected_shape=(num_instances, 1), expected_dtype=wp.transformf, @@ -648,7 +650,7 @@ def test_body_com_pose_b(self, backend, num_instances, device, rigid_object_ifac def test_body_mass(self, backend, num_instances, device, rigid_object_iface): obj, _ = rigid_object_iface obj.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( obj.data.body_mass, expected_shape=(num_instances, 1), expected_dtype=wp.float32, name="body_mass" ) @@ -658,7 +660,7 @@ def test_body_mass(self, backend, num_instances, device, rigid_object_iface): def test_body_inertia(self, backend, num_instances, device, rigid_object_iface): obj, _ = rigid_object_iface obj.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( obj.data.body_inertia, expected_shape=(num_instances, 1, 9), expected_dtype=wp.float32, name="body_inertia" ) @@ -668,7 +670,7 @@ def test_body_inertia(self, backend, num_instances, device, rigid_object_iface): def test_body_link_pos_w(self, backend, num_instances, device, rigid_object_iface): obj, _ = rigid_object_iface obj.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( obj.data.body_link_pos_w, expected_shape=(num_instances, 1), expected_dtype=wp.vec3f, name="body_link_pos_w" ) @@ -678,7 +680,7 @@ def test_body_link_pos_w(self, backend, num_instances, device, rigid_object_ifac def test_body_link_quat_w(self, backend, num_instances, device, rigid_object_iface): obj, _ = rigid_object_iface obj.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( obj.data.body_link_quat_w, expected_shape=(num_instances, 1), expected_dtype=wp.quatf, @@ -691,7 +693,7 @@ def test_body_link_quat_w(self, backend, num_instances, device, rigid_object_ifa def test_body_link_lin_vel_w(self, backend, num_instances, device, rigid_object_iface): obj, _ = rigid_object_iface obj.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( obj.data.body_link_lin_vel_w, expected_shape=(num_instances, 1), expected_dtype=wp.vec3f, @@ -704,7 +706,7 @@ def test_body_link_lin_vel_w(self, backend, num_instances, device, rigid_object_ def test_body_link_ang_vel_w(self, backend, num_instances, device, rigid_object_iface): obj, _ = rigid_object_iface obj.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( obj.data.body_link_ang_vel_w, expected_shape=(num_instances, 1), expected_dtype=wp.vec3f, @@ -717,7 +719,7 @@ def test_body_link_ang_vel_w(self, backend, num_instances, device, rigid_object_ def test_body_com_pos_w(self, backend, num_instances, device, rigid_object_iface): obj, _ = rigid_object_iface obj.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( obj.data.body_com_pos_w, expected_shape=(num_instances, 1), expected_dtype=wp.vec3f, name="body_com_pos_w" ) @@ -727,7 +729,7 @@ def test_body_com_pos_w(self, backend, num_instances, device, rigid_object_iface def test_body_com_quat_w(self, backend, num_instances, device, rigid_object_iface): obj, _ = rigid_object_iface obj.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( obj.data.body_com_quat_w, expected_shape=(num_instances, 1), expected_dtype=wp.quatf, name="body_com_quat_w" ) @@ -737,7 +739,7 @@ def test_body_com_quat_w(self, backend, num_instances, device, rigid_object_ifac def test_body_com_pos_b(self, backend, num_instances, device, rigid_object_iface): obj, _ = rigid_object_iface obj.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( obj.data.body_com_pos_b, expected_shape=(num_instances, 1), expected_dtype=wp.vec3f, name="body_com_pos_b" ) @@ -747,7 +749,7 @@ def test_body_com_pos_b(self, backend, num_instances, device, rigid_object_iface def test_body_com_quat_b(self, backend, num_instances, device, rigid_object_iface): obj, _ = rigid_object_iface obj.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( obj.data.body_com_quat_b, expected_shape=(num_instances, 1), expected_dtype=wp.quatf, name="body_com_quat_b" ) @@ -766,7 +768,7 @@ class TestRigidObjectDataDefaults: def test_default_root_pose(self, backend, num_instances, device, rigid_object_iface): obj, _ = rigid_object_iface obj.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( obj.data.default_root_pose, expected_shape=(num_instances,), expected_dtype=wp.transformf, @@ -779,7 +781,7 @@ def test_default_root_pose(self, backend, num_instances, device, rigid_object_if def test_default_root_vel(self, backend, num_instances, device, rigid_object_iface): obj, _ = rigid_object_iface obj.data.update(dt=0.01) - _check_wp_array( + _check_proxy_array( obj.data.default_root_vel, expected_shape=(num_instances,), expected_dtype=wp.spatial_vectorf, diff --git a/source/isaaclab/test/controllers/test_differential_ik.py b/source/isaaclab/test/controllers/test_differential_ik.py index a56545d31fd8..47abc78e0de6 100644 --- a/source/isaaclab/test/controllers/test_differential_ik.py +++ b/source/isaaclab/test/controllers/test_differential_ik.py @@ -152,8 +152,8 @@ def _run_ik_controller( ee_pose_b_des = torch.zeros(num_envs, diff_ik_controller.action_dim, device=sim.device) ee_pose_b_des[:] = ee_pose_b_des_set[current_goal_idx] # Compute current pose of the end-effector - ee_pose_w = wp.to_torch(robot.data.body_pose_w)[:, ee_frame_idx] - root_pose_w = wp.to_torch(robot.data.root_pose_w) + ee_pose_w = robot.data.body_pose_w.torch[:, ee_frame_idx] + root_pose_w = robot.data.root_pose_w.torch ee_pos_b, ee_quat_b = subtract_frame_transforms( root_pose_w[:, 0:3], root_pose_w[:, 3:7], ee_pose_w[:, 0:3], ee_pose_w[:, 3:7] ) @@ -175,14 +175,14 @@ def _run_ik_controller( torch.testing.assert_close(pos_error_norm, des_error, rtol=0.0, atol=1e-3) torch.testing.assert_close(rot_error_norm, des_error, rtol=0.0, atol=1e-3) # reset joint state - joint_pos = wp.to_torch(robot.data.default_joint_pos).clone() - joint_vel = wp.to_torch(robot.data.default_joint_vel).clone() + joint_pos = robot.data.default_joint_pos.torch.clone() + joint_vel = robot.data.default_joint_vel.torch.clone() # joint_pos *= sample_uniform(0.9, 1.1, joint_pos.shape, joint_pos.device) robot.write_joint_state_to_sim(joint_pos, joint_vel) robot.set_joint_position_target(joint_pos) robot.write_data_to_sim() # randomize root state yaw, ik should work regardless base rotation - root_state = wp.to_torch(robot.data.root_state_w).clone() + root_state = robot.data.root_state_w.torch.clone() root_state[:, 3:7] = random_yaw_orientation(num_envs, sim.device) robot.write_root_pose_to_sim(root_state[:, :7]) robot.write_root_velocity_to_sim(root_state[:, 7:]) @@ -200,13 +200,13 @@ def _run_ik_controller( # so we MUST skip the first step # obtain quantities from simulation jacobian = wp.to_torch(robot.root_view.get_jacobians())[:, ee_jacobi_idx, :, arm_joint_ids] - ee_pose_w = wp.to_torch(robot.data.body_pose_w)[:, ee_frame_idx] - root_pose_w = wp.to_torch(robot.data.root_pose_w) + ee_pose_w = robot.data.body_pose_w.torch[:, ee_frame_idx] + root_pose_w = robot.data.root_pose_w.torch base_rot = root_pose_w[:, 3:7] base_rot_matrix = matrix_from_quat(quat_inv(base_rot)) jacobian[:, :3, :] = torch.bmm(base_rot_matrix, jacobian[:, :3, :]) jacobian[:, 3:, :] = torch.bmm(base_rot_matrix, jacobian[:, 3:, :]) - joint_pos = wp.to_torch(robot.data.joint_pos)[:, arm_joint_ids] + joint_pos = robot.data.joint_pos.torch[:, arm_joint_ids] # compute frame in root frame ee_pos_b, ee_quat_b = subtract_frame_transforms( root_pose_w[:, 0:3], root_pose_w[:, 3:7], ee_pose_w[:, 0:3], ee_pose_w[:, 3:7] diff --git a/source/isaaclab/test/controllers/test_operational_space.py b/source/isaaclab/test/controllers/test_operational_space.py index 49a7efd425d3..63c1633c5f7c 100644 --- a/source/isaaclab/test/controllers/test_operational_space.py +++ b/source/isaaclab/test/controllers/test_operational_space.py @@ -1456,7 +1456,7 @@ def _run_op_space_controller( robot.update(dt=sim_dt) # Get the center of the robot soft joint limits - joint_centers = torch.mean(wp.to_torch(robot.data.soft_joint_pos_limits)[:, arm_joint_ids, :], dim=-1) + joint_centers = torch.mean(robot.data.soft_joint_pos_limits.torch[:, arm_joint_ids, :], dim=-1) # get the updated states ( @@ -1496,8 +1496,8 @@ def _run_op_space_controller( osc, ee_pose_b, ee_target_pose_b, ee_force_b, command, pos_mask, rot_mask, force_mask, frame ) # reset joint state to default - default_joint_pos = wp.to_torch(robot.data.default_joint_pos).clone() - default_joint_vel = wp.to_torch(robot.data.default_joint_vel).clone() + default_joint_pos = robot.data.default_joint_pos.torch.clone() + default_joint_vel = robot.data.default_joint_vel.torch.clone() robot.write_joint_state_to_sim(default_joint_pos, default_joint_vel) robot.set_joint_effort_target(zero_joint_efforts) # Set zero torques in the initial step robot.write_data_to_sim() @@ -1595,28 +1595,24 @@ def _update_states( gravity = wp.to_torch(robot.root_view.get_gravity_compensation_forces())[:, arm_joint_ids] # Convert the Jacobian from world to root frame jacobian_b = jacobian_w.clone() - root_rot_matrix = matrix_from_quat(quat_inv(wp.to_torch(robot.data.root_quat_w))) + root_rot_matrix = matrix_from_quat(quat_inv(robot.data.root_quat_w.torch)) jacobian_b[:, :3, :] = torch.bmm(root_rot_matrix, jacobian_b[:, :3, :]) jacobian_b[:, 3:, :] = torch.bmm(root_rot_matrix, jacobian_b[:, 3:, :]) # Compute current pose of the end-effector - root_pose_w = wp.to_torch(robot.data.root_pose_w) - ee_pose_w = wp.to_torch(robot.data.body_pose_w)[:, ee_frame_idx] + root_pose_w = robot.data.root_pose_w.torch + ee_pose_w = robot.data.body_pose_w.torch[:, ee_frame_idx] ee_pos_b, ee_quat_b = subtract_frame_transforms( root_pose_w[:, 0:3], root_pose_w[:, 3:7], ee_pose_w[:, 0:3], ee_pose_w[:, 3:7] ) ee_pose_b = torch.cat([ee_pos_b, ee_quat_b], dim=-1) # Compute the current velocity of the end-effector - ee_vel_w = wp.to_torch(robot.data.body_vel_w)[ - :, ee_frame_idx, : - ] # Extract end-effector velocity in the world frame - root_vel_w = wp.to_torch(robot.data.root_vel_w) # Extract root velocity in the world frame + ee_vel_w = robot.data.body_vel_w.torch[:, ee_frame_idx, :] # Extract end-effector velocity in the world frame + root_vel_w = robot.data.root_vel_w.torch # Extract root velocity in the world frame relative_vel_w = ee_vel_w - root_vel_w # Compute the relative velocity in the world frame - ee_lin_vel_b = quat_apply_inverse( - wp.to_torch(robot.data.root_quat_w), relative_vel_w[:, 0:3] - ) # From world to root frame - ee_ang_vel_b = quat_apply_inverse(wp.to_torch(robot.data.root_quat_w), relative_vel_w[:, 3:6]) + ee_lin_vel_b = quat_apply_inverse(robot.data.root_quat_w.torch, relative_vel_w[:, 0:3]) # From world to root frame + ee_ang_vel_b = quat_apply_inverse(robot.data.root_quat_w.torch, relative_vel_w[:, 3:6]) ee_vel_b = torch.cat([ee_lin_vel_b, ee_ang_vel_b], dim=-1) # Calculate the contact force @@ -1626,14 +1622,14 @@ def _update_states( contact_forces.update(sim_dt) # update contact sensor # Calculate the contact force by averaging over last four time steps (i.e., to smoothen) and # taking the max of three surfaces as only one should be the contact of interest - ee_force_w, _ = torch.max(torch.mean(wp.to_torch(contact_forces.data.net_forces_w_history), dim=1), dim=1) + ee_force_w, _ = torch.max(torch.mean(contact_forces.data.net_forces_w_history.torch, dim=1), dim=1) # This is a simplification, only for the sake of testing. ee_force_b = ee_force_w # Get joint positions and velocities - joint_pos = wp.to_torch(robot.data.joint_pos)[:, arm_joint_ids] - joint_vel = wp.to_torch(robot.data.joint_vel)[:, arm_joint_ids] + joint_pos = robot.data.joint_pos.torch[:, arm_joint_ids] + joint_vel = robot.data.joint_vel.torch[:, arm_joint_ids] return ( jacobian_b, diff --git a/source/isaaclab/test/controllers/test_pink_ik.py b/source/isaaclab/test/controllers/test_pink_ik.py index 8a6ce4927c34..4c7669b1c72e 100644 --- a/source/isaaclab/test/controllers/test_pink_ik.py +++ b/source/isaaclab/test/controllers/test_pink_ik.py @@ -21,7 +21,6 @@ import numpy as np import pytest import torch -import warp as wp from pink.configuration import Configuration from pink.tasks import FrameTask @@ -290,7 +289,7 @@ def run_movement_test(test_setup, test_config, test_cfg, aux_function=None): def get_link_pose(env, link_name): """Get the position and orientation of a link.""" link_index = env.scene["robot"].data.body_names.index(link_name) - link_states = wp.to_torch(env.scene._articulations["robot"].data.body_link_state_w) + link_states = env.scene._articulations["robot"].data.body_link_state_w.torch link_pose = link_states[:, link_index, :7] return link_pose[:, :3], link_pose[:, 3:7] @@ -343,7 +342,7 @@ def compute_errors( isaaclab_controlled_joint_ids = action_term._isaaclab_controlled_joint_ids # Get current and target positions for controlled joints only - curr_joints = wp.to_torch(articulation.data.joint_pos)[:, isaaclab_controlled_joint_ids].cpu().numpy()[0] + curr_joints = articulation.data.joint_pos.torch[:, isaaclab_controlled_joint_ids].cpu().numpy()[0] target_joints = action_term.processed_actions[:, : len(isaaclab_controlled_joint_ids)].cpu().numpy()[0] # Reorder joints for Pink IK (using controlled joint ordering) diff --git a/source/isaaclab/test/envs/check_manager_based_env_floating_cube.py b/source/isaaclab/test/envs/check_manager_based_env_floating_cube.py index 06123ad07a63..4762cc855cc5 100644 --- a/source/isaaclab/test/envs/check_manager_based_env_floating_cube.py +++ b/source/isaaclab/test/envs/check_manager_based_env_floating_cube.py @@ -31,7 +31,6 @@ """Rest everything follows.""" import torch -import warp as wp import isaaclab.envs.mdp as mdp import isaaclab.sim as sim_utils @@ -129,8 +128,8 @@ def process_actions(self, actions: torch.Tensor): def apply_actions(self): # implement a PD controller to track the target position - pos_error = self._processed_actions - (wp.to_torch(self._asset.data.root_pos_w) - self._env.scene.env_origins) - vel_error = -wp.to_torch(self._asset.data.root_lin_vel_w) + pos_error = self._processed_actions - (self._asset.data.root_pos_w.torch - self._env.scene.env_origins) + vel_error = -self._asset.data.root_lin_vel_w.torch # set velocity targets self._vel_command[:, :3] = self.p_gain * pos_error + self.d_gain * vel_error self._asset.write_root_velocity_to_sim(self._vel_command) @@ -152,7 +151,7 @@ def base_position(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg) -> torch.Tens """Root linear velocity in the asset's root frame.""" # extract the used quantities (to enable type-hinting) asset: RigidObject = env.scene[asset_cfg.name] - return wp.to_torch(asset.data.root_pos_w) - env.scene.env_origins + return asset.data.root_pos_w.torch - env.scene.env_origins ## diff --git a/source/isaaclab/test/envs/test_modify_env_param_curr_term.py b/source/isaaclab/test/envs/test_modify_env_param_curr_term.py index 9f58cbabc5bb..683b6eae02c7 100644 --- a/source/isaaclab/test/envs/test_modify_env_param_curr_term.py +++ b/source/isaaclab/test/envs/test_modify_env_param_curr_term.py @@ -12,7 +12,6 @@ import pytest import torch -import warp as wp import isaaclab.envs.mdp as mdp import isaaclab.sim as sim_utils @@ -89,14 +88,14 @@ def test_curriculum_modify_env_param(device): # test before curriculum kicks in, value agrees with default configuration joint_ids = env.event_manager.cfg.reset_cart_position.params["asset_cfg"].joint_ids assert env.observation_manager.cfg.policy.joint_pos_rel.func == mdp.joint_pos_rel - assert torch.any(wp.to_torch(robot.data.joint_pos)[:, joint_ids] != 0.0) + assert torch.any(robot.data.joint_pos.torch[:, joint_ids] != 0.0) assert env.max_episode_length_s == env_cfg.episode_length_s if count == 2: # test after curriculum makes effect, value agrees with new values assert env.observation_manager.cfg.policy.joint_pos_rel.func == mdp.joint_pos joint_ids = env.event_manager.cfg.reset_cart_position.params["asset_cfg"].joint_ids - assert torch.all(wp.to_torch(robot.data.joint_pos)[:, joint_ids] == 0.0) + assert torch.all(robot.data.joint_pos.torch[:, joint_ids] == 0.0) assert env.max_episode_length_s == 20 env.step(actions) diff --git a/source/isaaclab/test/envs/test_scale_randomization.py b/source/isaaclab/test/envs/test_scale_randomization.py index e3614e260a16..af5dc220e63f 100644 --- a/source/isaaclab/test/envs/test_scale_randomization.py +++ b/source/isaaclab/test/envs/test_scale_randomization.py @@ -22,7 +22,6 @@ import pytest import torch -import warp as wp from pxr import Sdf @@ -102,8 +101,8 @@ def process_actions(self, actions: torch.Tensor): def apply_actions(self): # implement a PD controller to track the target position - pos_error = self._processed_actions - (wp.to_torch(self._asset.data.root_pos_w) - self._env.scene.env_origins) - vel_error = -wp.to_torch(self._asset.data.root_lin_vel_w) + pos_error = self._processed_actions - (self._asset.data.root_pos_w.torch - self._env.scene.env_origins) + vel_error = -self._asset.data.root_lin_vel_w.torch # set velocity targets self._vel_command[:, :3] = self.p_gain * pos_error + self.d_gain * vel_error self._asset.write_root_velocity_to_sim(self._vel_command) @@ -131,7 +130,7 @@ def base_position(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg) -> torch.Tens """Root linear velocity in the asset's root frame.""" # extract the used quantities (to enable type-hinting) asset: RigidObject = env.scene[asset_cfg.name] - return wp.to_torch(asset.data.root_pos_w) - env.scene.env_origins + return asset.data.root_pos_w.torch - env.scene.env_origins ## diff --git a/source/isaaclab/test/markers/check_markers_visibility.py b/source/isaaclab/test/markers/check_markers_visibility.py index 98dbee8ddcd7..eadc8af82a53 100644 --- a/source/isaaclab/test/markers/check_markers_visibility.py +++ b/source/isaaclab/test/markers/check_markers_visibility.py @@ -96,14 +96,14 @@ def run_simulator( count = 0 # reset the scene entities # root state - root_state = scene["robot"].data.default_root_state.clone() + root_state = scene["robot"].data.default_root_state.torch.clone() root_state[:, :3] += scene.env_origins scene["robot"].write_root_pose_to_sim(root_state[:, :7]) scene["robot"].write_root_velocity_to_sim(root_state[:, 7:]) # set joint positions with some noise joint_pos, joint_vel = ( - scene["robot"].data.default_joint_pos.clone(), - scene["robot"].data.default_joint_vel.clone(), + scene["robot"].data.default_joint_pos.torch.clone(), + scene["robot"].data.default_joint_vel.torch.clone(), ) scene["robot"].write_joint_state_to_sim(joint_pos, joint_vel) # clear internal buffers diff --git a/source/isaaclab/test/scene/check_interactive_scene.py b/source/isaaclab/test/scene/check_interactive_scene.py index 5b2463b315a9..df0c4a12a945 100644 --- a/source/isaaclab/test/scene/check_interactive_scene.py +++ b/source/isaaclab/test/scene/check_interactive_scene.py @@ -103,8 +103,8 @@ def main(): print("[INFO]: Setup complete...") # default joint targets - robot_1_actions = scene.articulations["robot_1"].data.default_joint_pos.clone() - robot_2_actions = scene.articulations["robot_2"].data.default_joint_pos.clone() + robot_1_actions = scene.articulations["robot_1"].data.default_joint_pos.torch.clone() + robot_2_actions = scene.articulations["robot_2"].data.default_joint_pos.torch.clone() # Define simulation stepping sim_dt = sim.get_physics_dt() sim_time = 0.0 @@ -124,10 +124,10 @@ def main(): sim_time = 0.0 count = 0 # reset root state - root_state = scene.articulations["robot_1"].data.default_root_state.clone() + root_state = scene.articulations["robot_1"].data.default_root_state.torch.clone() root_state[:, :3] += scene.env_origins - joint_pos = scene.articulations["robot_1"].data.default_joint_pos - joint_vel = scene.articulations["robot_1"].data.default_joint_vel + joint_pos = scene.articulations["robot_1"].data.default_joint_pos.torch + joint_vel = scene.articulations["robot_1"].data.default_joint_vel.torch # -- set root state # -- robot 1 scene.articulations["robot_1"].write_root_pose_to_sim(root_state[:, :7]) diff --git a/source/isaaclab/test/scene/test_interactive_scene.py b/source/isaaclab/test/scene/test_interactive_scene.py index 20b96aaacd4f..5c6954080149 100644 --- a/source/isaaclab/test/scene/test_interactive_scene.py +++ b/source/isaaclab/test/scene/test_interactive_scene.py @@ -16,7 +16,6 @@ import pytest import torch -import warp as wp import isaaclab.sim as sim_utils from isaaclab.actuators import ImplicitActuatorCfg @@ -84,8 +83,8 @@ def test_relative_flag(device, setup_scene): # test is relative == False prev_state = scene.get_state(is_relative=False) scene["robot"].write_joint_state_to_sim( - position=torch.rand_like(wp.to_torch(scene["robot"].data.joint_pos)), - velocity=torch.rand_like(wp.to_torch(scene["robot"].data.joint_pos)), + position=torch.rand_like(scene["robot"].data.joint_pos.torch), + velocity=torch.rand_like(scene["robot"].data.joint_pos.torch), ) next_state = scene.get_state(is_relative=False) assert_state_different(prev_state, next_state) @@ -95,8 +94,8 @@ def test_relative_flag(device, setup_scene): # test is relative == True prev_state = scene.get_state(is_relative=True) scene["robot"].write_joint_state_to_sim( - position=torch.rand_like(wp.to_torch(scene["robot"].data.joint_pos)), - velocity=torch.rand_like(wp.to_torch(scene["robot"].data.joint_pos)), + position=torch.rand_like(scene["robot"].data.joint_pos.torch), + velocity=torch.rand_like(scene["robot"].data.joint_pos.torch), ) next_state = scene.get_state(is_relative=True) assert_state_different(prev_state, next_state) @@ -114,16 +113,16 @@ def test_reset_to_env_ids_input_types(device, setup_scene): # test env_ids = None prev_state = scene.get_state() scene["robot"].write_joint_state_to_sim( - position=torch.rand_like(wp.to_torch(scene["robot"].data.joint_pos)), - velocity=torch.rand_like(wp.to_torch(scene["robot"].data.joint_pos)), + position=torch.rand_like(scene["robot"].data.joint_pos.torch), + velocity=torch.rand_like(scene["robot"].data.joint_pos.torch), ) scene.reset_to(prev_state, env_ids=None) assert_state_equal(prev_state, scene.get_state()) # test env_ids = torch tensor scene["robot"].write_joint_state_to_sim( - position=torch.rand_like(wp.to_torch(scene["robot"].data.joint_pos)), - velocity=torch.rand_like(wp.to_torch(scene["robot"].data.joint_pos)), + position=torch.rand_like(scene["robot"].data.joint_pos.torch), + velocity=torch.rand_like(scene["robot"].data.joint_pos.torch), ) scene.reset_to(prev_state, env_ids=torch.arange(scene.num_envs, device=scene.device, dtype=torch.int32)) assert_state_equal(prev_state, scene.get_state()) diff --git a/source/isaaclab/test/sensors/check_multi_mesh_ray_caster.py b/source/isaaclab/test/sensors/check_multi_mesh_ray_caster.py index 2ffd89edf2f4..eed776431279 100644 --- a/source/isaaclab/test/sensors/check_multi_mesh_ray_caster.py +++ b/source/isaaclab/test/sensors/check_multi_mesh_ray_caster.py @@ -164,8 +164,8 @@ def main(): print(ray_caster) # Get the initial positions of the balls - ball_initial_poses = balls.data.root_pose_w.clone() - ball_initial_velocities = balls.data.root_vel_w.clone() + ball_initial_poses = balls.data.root_pose_w.torch.clone() + ball_initial_velocities = balls.data.root_vel_w.torch.clone() # Create a counter for resetting the scene step_count = 0 diff --git a/source/isaaclab/test/sensors/check_ray_caster.py b/source/isaaclab/test/sensors/check_ray_caster.py index dc815afc1464..821fc247743e 100644 --- a/source/isaaclab/test/sensors/check_ray_caster.py +++ b/source/isaaclab/test/sensors/check_ray_caster.py @@ -133,8 +133,8 @@ def main(): print(ray_caster) # Get the initial positions of the balls - ball_initial_poses = balls.data.root_pose_w.clone() - ball_initial_velocities = balls.data.root_vel_w.clone() + ball_initial_poses = balls.data.root_pose_w.torch.clone() + ball_initial_velocities = balls.data.root_vel_w.torch.clone() # Create a counter for resetting the scene step_count = 0 diff --git a/source/isaaclab/test/sensors/test_ray_caster.py b/source/isaaclab/test/sensors/test_ray_caster.py index 4e29b25ce351..6ac63cf0d041 100644 --- a/source/isaaclab/test/sensors/test_ray_caster.py +++ b/source/isaaclab/test/sensors/test_ray_caster.py @@ -291,9 +291,8 @@ def test_raycaster_offset_does_not_affect_pos_w(): sim.reset() sensor.update(dt) - # data.pos_w / data.ray_hits_w are wp.array after the ray caster warp-backend - # migration (PR #4967); convert to torch views for indexing. - pos_w = wp.to_torch(sensor.data.pos_w)[0].cpu() + # data.pos_w / data.ray_hits_w return ProxyArray wrappers; use .torch for tensor indexing. + pos_w = sensor.data.pos_w.torch[0].cpu() # pos_w.z should be near the body height, NOT body_height + offset assert abs(pos_w[2].item() - body_pos[2]) < 1.0, ( @@ -302,7 +301,7 @@ def test_raycaster_offset_does_not_affect_pos_w(): ) # ray_hits should be near z=0 (ground plane) - hits_z = wp.to_torch(sensor.data.ray_hits_w)[0, :, 2].cpu() + hits_z = sensor.data.ray_hits_w.torch[0, :, 2].cpu() valid = hits_z[~torch.isinf(hits_z)] if len(valid) > 0: assert valid.abs().max().item() < 2.0, ( diff --git a/source/isaaclab/test/sensors/test_ray_caster_integration.py b/source/isaaclab/test/sensors/test_ray_caster_integration.py index 62b10a679661..36b6d8ff4692 100644 --- a/source/isaaclab/test/sensors/test_ray_caster_integration.py +++ b/source/isaaclab/test/sensors/test_ray_caster_integration.py @@ -120,7 +120,7 @@ def test_articulation_view_path(sim_ground): sim.reset() sensor.update(_DT) - pos_w = wp.to_torch(sensor.data.pos_w)[0].cpu().numpy() + pos_w = sensor.data.pos_w.torch[0].cpu().numpy() np.testing.assert_allclose( pos_w, expected_pos, @@ -128,7 +128,7 @@ def test_articulation_view_path(sim_ground): err_msg="ArticulationView: sensor pos_w must match initial prim position", ) - hits = wp.to_torch(sensor.data.ray_hits_w)[0, 0].cpu().numpy() + hits = sensor.data.ray_hits_w.torch[0, 0].cpu().numpy() assert abs(hits[2]) < 0.5, f"ArticulationView: downward ray should hit near z=0, got z={hits[2]}" @@ -164,7 +164,7 @@ def test_rigid_body_view_path(sim_ground): sim.reset() sensor.update(_DT) - pos_w = wp.to_torch(sensor.data.pos_w)[0].cpu().numpy() + pos_w = sensor.data.pos_w.torch[0].cpu().numpy() np.testing.assert_allclose( pos_w, expected_pos, @@ -172,7 +172,7 @@ def test_rigid_body_view_path(sim_ground): err_msg="RigidBodyView: sensor pos_w must match initial prim position", ) - hits = wp.to_torch(sensor.data.ray_hits_w)[0, 0].cpu().numpy() + hits = sensor.data.ray_hits_w.torch[0, 0].cpu().numpy() assert abs(hits[2]) < 0.5, f"RigidBodyView: downward ray should hit near z=0, got z={hits[2]}" @@ -334,13 +334,13 @@ def test_multi_mesh_env_mask_preserves_masked_buffers(sim_ground): # First update: populate buffers with real values sensor.update(_DT) - hits_before = wp.to_torch(sensor.data.ray_hits_w).clone() + hits_before = sensor.data.ray_hits_w.torch.clone() # Second update with env masked out: buffers must not change mask_all_false = wp.array([False], dtype=wp.bool, device=sensor.device) sensor._update_buffers_impl(mask_all_false) - hits_after = wp.to_torch(sensor.data.ray_hits_w) + hits_after = sensor.data.ray_hits_w.torch torch.testing.assert_close( hits_after, hits_before, diff --git a/source/isaaclab/test/sensors/test_ray_caster_sensor.py b/source/isaaclab/test/sensors/test_ray_caster_sensor.py index f1c90a986b5d..3326c524f958 100644 --- a/source/isaaclab/test/sensors/test_ray_caster_sensor.py +++ b/source/isaaclab/test/sensors/test_ray_caster_sensor.py @@ -14,7 +14,6 @@ import numpy as np import pytest import torch -import warp as wp import isaaclab.sim as sim_utils from isaaclab.sensors.ray_caster import RayCaster, RayCasterCfg, patterns @@ -97,9 +96,9 @@ def test_world_alignment_ignores_sensor_pitch(sim_ground): sensor_upright.update(dt) sensor_pitched.update(dt) - # ray_hits_w is a wp.array(dtype=wp.vec3f); convert to torch for indexing - hits_upright = wp.to_torch(sensor_upright.data.ray_hits_w) # (1, 1, 3) - hits_pitched = wp.to_torch(sensor_pitched.data.ray_hits_w) + # ray_hits_w returns a ProxyArray; use .torch for tensor indexing. + hits_upright = sensor_upright.data.ray_hits_w.torch # (1, 1, 3) + hits_pitched = sensor_pitched.data.ray_hits_w.torch # Both must hit z=0 (straight down, world frame direction) assert abs(hits_upright[0, 0, 2].item()) < 0.02, ( @@ -139,8 +138,8 @@ def test_base_alignment_rotates_ray_direction(sim_ground): sensor_world.update(dt) sensor_base.update(dt) - hits_world = wp.to_torch(sensor_world.data.ray_hits_w) # (1, 1, 3) - hits_base = wp.to_torch(sensor_base.data.ray_hits_w) + hits_world = sensor_world.data.ray_hits_w.torch # (1, 1, 3) + hits_base = sensor_base.data.ray_hits_w.torch # World mode: ray still hits directly below (x≈0, y≈0, z≈0) assert abs(hits_world[0, 0, 0].item()) < 0.05, f"World mode hit x must be near 0, got {hits_world[0, 0, 0].item()}" @@ -201,8 +200,8 @@ def _cfg_with_offset(prim_path, alignment): sensor_world.update(dt) sensor_yaw.update(dt) - hits_world = wp.to_torch(sensor_world.data.ray_hits_w) # (1, 1, 3) - hits_yaw = wp.to_torch(sensor_yaw.data.ray_hits_w) + hits_world = sensor_world.data.ray_hits_w.torch # (1, 1, 3) + hits_yaw = sensor_yaw.data.ray_hits_w.torch # Both modes must hit the ground (direction unchanged = straight down in both modes) assert abs(hits_world[0, 0, 2].item()) < 0.05, "World mode must hit z≈0" diff --git a/source/isaaclab/test/sim/frame_view_contract_utils.py b/source/isaaclab/test/sim/frame_view_contract_utils.py index 37734fee4322..4336f94bd3d7 100644 --- a/source/isaaclab/test/sim/frame_view_contract_utils.py +++ b/source/isaaclab/test/sim/frame_view_contract_utils.py @@ -43,6 +43,8 @@ class ViewBundle(NamedTuple): import torch import warp as wp +from isaaclab.utils.warp import ProxyArray + CHILD_OFFSET = (0.1, 0.0, 0.05) """Local offset of the child prim from its parent, shared by all backend fixtures.""" @@ -60,7 +62,9 @@ class ViewBundle(NamedTuple): def _t(a): - """Convert wp.array to torch.Tensor (pass-through for Tensor).""" + """Convert a wp.array or ProxyArray return to a torch.Tensor (pass-through otherwise).""" + if isinstance(a, ProxyArray): + return a.torch return wp.to_torch(a) if isinstance(a, wp.array) else a @@ -357,3 +361,44 @@ def test_set_world_indexed_only_affects_subset(device, view_factory): torch.testing.assert_close(updated[3], _t(new_pos)[1], atol=ATOL, rtol=0) finally: bundle.teardown() + + +@pytest.mark.parametrize("device", ["cpu", "cuda:0"]) +def test_return_types_are_torcharray(device, view_factory): + """Public API contract — every backend returns ProxyArray from the pose getters.""" + bundle = view_factory(num_envs=2, device=device) + try: + pos_full, quat_full = bundle.view.get_world_poses() + assert isinstance(pos_full, ProxyArray), ( + f"get_world_poses()[0] must be ProxyArray, got {type(pos_full).__name__}" + ) + assert isinstance(quat_full, ProxyArray), ( + f"get_world_poses()[1] must be ProxyArray, got {type(quat_full).__name__}" + ) + + indices = wp.array([0], dtype=wp.int32, device=bundle.view.device) + pos_idx, quat_idx = bundle.view.get_world_poses(indices) + assert isinstance(pos_idx, ProxyArray), ( + f"get_world_poses(indices)[0] must be ProxyArray, got {type(pos_idx).__name__}" + ) + assert isinstance(quat_idx, ProxyArray), ( + f"get_world_poses(indices)[1] must be ProxyArray, got {type(quat_idx).__name__}" + ) + + lpos_full, lquat_full = bundle.view.get_local_poses() + assert isinstance(lpos_full, ProxyArray), ( + f"get_local_poses()[0] must be ProxyArray, got {type(lpos_full).__name__}" + ) + assert isinstance(lquat_full, ProxyArray), ( + f"get_local_poses()[1] must be ProxyArray, got {type(lquat_full).__name__}" + ) + + lpos_idx, lquat_idx = bundle.view.get_local_poses(indices) + assert isinstance(lpos_idx, ProxyArray), ( + f"get_local_poses(indices)[0] must be ProxyArray, got {type(lpos_idx).__name__}" + ) + assert isinstance(lquat_idx, ProxyArray), ( + f"get_local_poses(indices)[1] must be ProxyArray, got {type(lquat_idx).__name__}" + ) + finally: + bundle.teardown() diff --git a/source/isaaclab/test/sim/test_views_xform_prim.py b/source/isaaclab/test/sim/test_views_xform_prim.py index c0f988567ef2..64cd86a7466f 100644 --- a/source/isaaclab/test/sim/test_views_xform_prim.py +++ b/source/isaaclab/test/sim/test_views_xform_prim.py @@ -229,7 +229,7 @@ def test_nested_hierarchy_world_poses(device): frames_view.set_local_poses(translations=torch.tensor(frame_positions, device=device)) targets_view.set_local_poses(translations=torch.tensor(target_positions, device=device)) - world_pos = wp.to_torch(targets_view.get_world_poses()[0]) + world_pos = targets_view.get_world_poses()[0].torch expected = torch.tensor( [[f[j] + t[j] for j in range(3)] for f, t in zip(frame_positions, target_positions)], device=device, @@ -258,7 +258,7 @@ def test_compare_get_world_poses_with_isaacsim(): isaaclab_view = FrameView(pattern, device="cpu") isaacsim_view = _IsaacSimXformPrimView(pattern, reset_xform_properties=False) - isaaclab_pos = wp.to_torch(isaaclab_view.get_world_poses()[0]) + isaaclab_pos = isaaclab_view.get_world_poses()[0].torch isaacsim_pos, isaacsim_quat = isaacsim_view.get_world_poses() if not isinstance(isaacsim_pos, torch.Tensor): isaacsim_pos = torch.tensor(isaacsim_pos, dtype=torch.float32) @@ -286,12 +286,12 @@ def test_with_franka_robots(device): view = FrameView("/World/Franka_.*", device=device) assert view.count == 2 - positions = wp.to_torch(view.get_world_poses()[0]) + positions = view.get_world_poses()[0].torch torch.testing.assert_close(positions, torch.zeros(2, 3, device=device), atol=1e-5, rtol=0) new_pos = torch.tensor([[10.0, 10.0, 0.0], [-40.0, -40.0, 0.0]], device=device) new_quat = torch.tensor([[0.0, 0.0, 0.7071068, 0.7071068], [0.0, 0.0, -0.7071068, 0.7071068]], device=device) view.set_world_poses(positions=new_pos, orientations=new_quat) - ret_pos = wp.to_torch(view.get_world_poses()[0]) + ret_pos = view.get_world_poses()[0].torch torch.testing.assert_close(ret_pos, new_pos, atol=1e-5, rtol=0) diff --git a/source/isaaclab/test/test_mock_interfaces/test_mock_assets.py b/source/isaaclab/test/test_mock_interfaces/test_mock_assets.py index 61d310e4f20d..a10fe5150ff6 100644 --- a/source/isaaclab/test/test_mock_interfaces/test_mock_assets.py +++ b/source/isaaclab/test/test_mock_interfaces/test_mock_assets.py @@ -55,47 +55,39 @@ def test_joint_state_shapes(self, robot): def test_root_state_shapes(self, robot): """Test root state tensor shapes.""" - import warp as wp - # Link frame - pose is wp.transformf so shape is (4,) but converts to (4, 7) - assert wp.to_torch(robot.data.root_link_pose_w).shape == (4, 7) + assert robot.data.root_link_pose_w.torch.shape == (4, 7) # vel is wp.spatial_vectorf so shape is (4,) but converts to (4, 6) - assert wp.to_torch(robot.data.root_link_vel_w).shape == (4, 6) - assert wp.to_torch(robot.data.root_link_state_w).shape == (4, 13) + assert robot.data.root_link_vel_w.torch.shape == (4, 6) + assert robot.data.root_link_state_w.torch.shape == (4, 13) # Sliced properties - assert wp.to_torch(robot.data.root_link_pos_w).shape == (4, 3) - assert wp.to_torch(robot.data.root_link_quat_w).shape == (4, 4) - assert wp.to_torch(robot.data.root_link_lin_vel_w).shape == (4, 3) - assert wp.to_torch(robot.data.root_link_ang_vel_w).shape == (4, 3) + assert robot.data.root_link_pos_w.torch.shape == (4, 3) + assert robot.data.root_link_quat_w.torch.shape == (4, 4) + assert robot.data.root_link_lin_vel_w.torch.shape == (4, 3) + assert robot.data.root_link_ang_vel_w.torch.shape == (4, 3) def test_body_state_shapes(self, robot): """Test body state tensor shapes.""" - import warp as wp - # body_link_pose_w is wp.transformf so shape is (4, 13) but converts to (4, 13, 7) - assert wp.to_torch(robot.data.body_link_pose_w).shape == (4, 13, 7) + assert robot.data.body_link_pose_w.torch.shape == (4, 13, 7) # body_link_vel_w is wp.spatial_vectorf so shape is (4, 13) but converts to (4, 13, 6) - assert wp.to_torch(robot.data.body_link_vel_w).shape == (4, 13, 6) - assert wp.to_torch(robot.data.body_link_state_w).shape == (4, 13, 13) + assert robot.data.body_link_vel_w.torch.shape == (4, 13, 6) + assert robot.data.body_link_state_w.torch.shape == (4, 13, 13) def test_default_state_shapes(self, robot): """Test default state tensor shapes.""" - import warp as wp - # default_root_pose is wp.transformf so shape is (4,) but converts to (4, 7) - assert wp.to_torch(robot.data.default_root_pose).shape == (4, 7) + assert robot.data.default_root_pose.torch.shape == (4, 7) # default_root_vel is wp.spatial_vectorf so shape is (4,) but converts to (4, 6) - assert wp.to_torch(robot.data.default_root_vel).shape == (4, 6) - assert wp.to_torch(robot.data.default_root_state).shape == (4, 13) + assert robot.data.default_root_vel.torch.shape == (4, 6) + assert robot.data.default_root_state.torch.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.""" - import warp as wp - - quat = wp.to_torch(robot.data.root_link_quat_w) + quat = robot.data.root_link_quat_w.torch # XYZW format: x=y=z=0, w=1 expected = torch.zeros_like(quat) expected[:, 3] = 1.0 # Set w=1 @@ -103,23 +95,19 @@ def test_identity_quaternion_default(self, robot): def test_set_joint_pos(self, robot): """Test setting joint positions.""" - import warp as wp - joint_pos = torch.randn(4, 12) robot.data.set_joint_pos(joint_pos) - assert torch.allclose(wp.to_torch(robot.data.joint_pos), joint_pos) + assert torch.allclose(robot.data.joint_pos.torch, joint_pos) def test_set_mock_data_bulk(self, robot): """Test bulk data setter.""" - import warp as wp - 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(wp.to_torch(robot.data.joint_pos), joint_pos) - assert torch.allclose(wp.to_torch(robot.data.joint_vel), joint_vel) + assert torch.allclose(robot.data.joint_pos.torch, joint_pos) + assert torch.allclose(robot.data.joint_vel.torch, joint_vel) def test_find_joints(self): """Test joint finding by regex.""" @@ -163,17 +151,13 @@ def test_find_bodies(self): def test_set_joint_position_target(self, robot): """Test setting joint position targets.""" - import warp as wp - target = torch.randn(4, 12) robot.set_joint_position_target(target) - assert torch.allclose(wp.to_torch(robot.data.joint_pos_target), target) + assert torch.allclose(robot.data.joint_pos_target.torch, target) def test_joint_limits(self, robot): """Test joint limits.""" - import warp as wp - - limits = wp.to_torch(robot.data.joint_pos_limits) + limits = robot.data.joint_pos_limits.torch assert limits.shape == (4, 12, 2) # Default limits should be -inf to inf assert torch.all(limits[..., 0] == float("-inf")) @@ -201,22 +185,18 @@ def test_initialization(self, obj): def test_root_state_shapes(self, obj): """Test root state tensor shapes.""" - import warp as wp - # root_link_pose_w is wp.transformf so shape is (4,) but converts to (4, 7) - assert wp.to_torch(obj.data.root_link_pose_w).shape == (4, 7) + assert obj.data.root_link_pose_w.torch.shape == (4, 7) # root_link_vel_w is wp.spatial_vectorf so shape is (4,) but converts to (4, 6) - assert wp.to_torch(obj.data.root_link_vel_w).shape == (4, 6) - assert wp.to_torch(obj.data.root_link_state_w).shape == (4, 13) + assert obj.data.root_link_vel_w.torch.shape == (4, 6) + assert obj.data.root_link_state_w.torch.shape == (4, 13) def test_body_state_shapes(self, obj): """Test body state tensor shapes (single body).""" - import warp as wp - # body_link_pose_w is wp.transformf so shape is (4, 1) but converts to (4, 1, 7) - assert wp.to_torch(obj.data.body_link_pose_w).shape == (4, 1, 7) + assert obj.data.body_link_pose_w.torch.shape == (4, 1, 7) # body_link_vel_w is wp.spatial_vectorf so shape is (4, 1) but converts to (4, 1, 6) - assert wp.to_torch(obj.data.body_link_vel_w).shape == (4, 1, 6) + assert obj.data.body_link_vel_w.torch.shape == (4, 1, 6) def test_body_properties(self, obj): """Test body property shapes.""" @@ -248,13 +228,11 @@ def test_initialization(self, collection): def test_body_state_shapes(self, collection): """Test body state tensor shapes.""" - import warp as wp - # body_link_pose_w is wp.transformf so shape is (4, 5) but converts to (4, 5, 7) - assert wp.to_torch(collection.data.body_link_pose_w).shape == (4, 5, 7) + assert collection.data.body_link_pose_w.torch.shape == (4, 5, 7) # body_link_vel_w is wp.spatial_vectorf so shape is (4, 5) but converts to (4, 5, 6) - assert wp.to_torch(collection.data.body_link_vel_w).shape == (4, 5, 6) - assert wp.to_torch(collection.data.body_link_state_w).shape == (4, 5, 13) + assert collection.data.body_link_vel_w.torch.shape == (4, 5, 6) + assert collection.data.body_link_state_w.torch.shape == (4, 5, 13) def test_find_bodies_returns_mask(self, collection): """Test that find_bodies returns a mask tensor.""" @@ -301,17 +279,14 @@ def test_create_mock_articulation(self): def test_create_mock_rigid_object(self): """Test rigid object factory function.""" - import warp as wp - obj = create_mock_rigid_object(num_instances=3) assert obj.num_instances == 3 assert obj.num_bodies == 1 # root_link_pose_w is wp.transformf so shape is (3,) but converts to (3, 7) - assert wp.to_torch(obj.data.root_link_pose_w).shape == (3, 7) + assert obj.data.root_link_pose_w.torch.shape == (3, 7) def test_create_mock_rigid_object_collection(self): """Test rigid object collection factory function.""" - import warp as wp collection = create_mock_rigid_object_collection( num_instances=4, @@ -322,7 +297,7 @@ def test_create_mock_rigid_object_collection(self): assert collection.num_bodies == 6 assert collection.body_names == ["obj_0", "obj_1", "obj_2", "obj_3", "obj_4", "obj_5"] # body_link_pose_w is wp.transformf so shape is (4, 6) but converts to (4, 6, 7) - assert wp.to_torch(collection.data.body_link_pose_w).shape == (4, 6, 7) + assert collection.data.body_link_pose_w.torch.shape == (4, 6, 7) # ============================================================================== @@ -349,7 +324,6 @@ def test_basic_build(self): def test_with_default_positions(self): """Test setting default joint positions.""" - import warp as wp default_pos = [0.0, 0.5, -0.5] robot = ( @@ -360,11 +334,10 @@ def test_with_default_positions(self): ) expected = torch.tensor([default_pos, default_pos]) - assert torch.allclose(wp.to_torch(robot.data.joint_pos), expected) + assert torch.allclose(robot.data.joint_pos.torch, expected) def test_with_joint_limits(self): """Test setting joint limits.""" - import warp as wp robot = ( MockArticulationBuilder() @@ -374,6 +347,6 @@ def test_with_joint_limits(self): .build() ) - limits = wp.to_torch(robot.data.joint_pos_limits) + limits = robot.data.joint_pos_limits.torch 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 index be45798a7d85..c688975ca3eb 100644 --- a/source/isaaclab/test/test_mock_interfaces/test_mock_data_properties.py +++ b/source/isaaclab/test/test_mock_interfaces/test_mock_data_properties.py @@ -7,7 +7,6 @@ import pytest import torch -import warp as wp from isaaclab.test.mock_interfaces.assets import ( MockArticulationData, @@ -47,7 +46,7 @@ def data(self): 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 wp.to_torch(prop).shape == expected_shape, f"{property_name} has wrong shape" + assert prop.torch.shape == expected_shape, f"{property_name} has wrong shape" @pytest.mark.parametrize( "setter_name,property_name,shape", @@ -58,25 +57,23 @@ def test_property_shapes(self, data, property_name, expected_shape): ) def test_setters_update_properties(self, data, setter_name, property_name, shape): """Test that setters properly update the corresponding properties.""" - import warp as wp test_value = torch.randn(shape) setter = getattr(data, setter_name) setter(test_value) - result = wp.to_torch(getattr(data, property_name)) + result = getattr(data, property_name).torch 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.""" - import warp as wp ang_vel = torch.randn(4, 3) lin_acc = torch.randn(4, 3) data.set_mock_data(ang_vel_b=ang_vel, lin_acc_b=lin_acc) - assert torch.allclose(wp.to_torch(data.ang_vel_b), ang_vel) - assert torch.allclose(wp.to_torch(data.lin_acc_b), lin_acc) + assert torch.allclose(data.ang_vel_b.torch, ang_vel) + assert torch.allclose(data.lin_acc_b.torch, lin_acc) # ============================================================================== @@ -108,7 +105,7 @@ def data(self): 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 wp.to_torch(prop).shape == expected_shape, f"{property_name} has wrong shape" + assert prop.torch.shape == expected_shape, f"{property_name} has wrong shape" @pytest.mark.parametrize( "setter_name,property_name,shape", @@ -124,17 +121,15 @@ def test_property_shapes(self, data, property_name, expected_shape): ) def test_setters_update_properties(self, data, setter_name, property_name, shape): """Test that setters properly update the corresponding properties.""" - import warp as wp test_value = torch.randn(shape) setter = getattr(data, setter_name) setter(test_value) - result = wp.to_torch(getattr(data, property_name)) + result = getattr(data, property_name).torch 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.""" - import warp as wp lin_vel = torch.randn(4, 3) ang_vel = torch.randn(4, 3) @@ -142,23 +137,21 @@ def test_bulk_setter(self, data): data.set_mock_data(lin_vel_b=lin_vel, ang_vel_b=ang_vel, lin_acc_b=lin_acc) - assert torch.allclose(wp.to_torch(data.lin_vel_b), lin_vel) - assert torch.allclose(wp.to_torch(data.ang_vel_b), ang_vel) - assert torch.allclose(wp.to_torch(data.lin_acc_b), lin_acc) + assert torch.allclose(data.lin_vel_b.torch, lin_vel) + assert torch.allclose(data.ang_vel_b.torch, ang_vel) + assert torch.allclose(data.lin_acc_b.torch, lin_acc) def test_default_quaternion_is_identity(self, data): """Test that default quaternion is identity in XYZW format: (x, y, z, w) = (0, 0, 0, 1).""" - import warp as wp - quat = wp.to_torch(data.quat_w) + quat = data.quat_w.torch assert torch.allclose(quat[:, :3], torch.zeros(4, 3)) # xyz=0 assert torch.allclose(quat[:, 3], torch.ones(4)) # w=1 def test_default_gravity_points_down(self, data): """Test that default gravity points in -z direction.""" - import warp as wp - gravity = wp.to_torch(data.projected_gravity_b) + gravity = data.projected_gravity_b.torch expected = torch.tensor([[0, 0, -1]] * 4, dtype=torch.float32) assert torch.allclose(gravity, expected) @@ -229,7 +222,7 @@ def test_optional_property_shapes_with_history(self, data, property_name, expect if property_name == "force_matrix_w_history": assert prop.shape == expected_shape, f"{property_name} has wrong shape" else: - assert wp.to_torch(prop).shape == expected_shape, f"{property_name} has wrong shape" + assert prop.torch.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.""" @@ -253,12 +246,11 @@ def test_optional_properties_none_without_config(self, data_no_history): ) def test_setters_update_properties(self, data, setter_name, property_name, shape): """Test that setters properly update the corresponding properties.""" - import warp as wp test_value = torch.randn(shape) setter = getattr(data, setter_name) setter(test_value) - result = wp.to_torch(getattr(data, property_name)) + result = getattr(data, property_name).torch assert torch.allclose(result, test_value), f"{setter_name} did not update {property_name}" @@ -297,7 +289,7 @@ def data(self): 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 wp.to_torch(prop).shape == expected_shape, f"{property_name} has wrong shape" + assert prop.torch.shape == expected_shape, f"{property_name} has wrong shape" @pytest.mark.parametrize( "setter_name,property_name,shape", @@ -312,12 +304,11 @@ def test_property_shapes(self, data, property_name, expected_shape): ) def test_setters_update_properties(self, data, setter_name, property_name, shape): """Test that setters properly update the corresponding properties.""" - import warp as wp test_value = torch.randn(shape) setter = getattr(data, setter_name) setter(test_value) - result = wp.to_torch(getattr(data, property_name)) + result = getattr(data, property_name).torch assert torch.allclose(result, test_value), f"{setter_name} did not update {property_name}" def test_target_frame_names(self, data): @@ -362,7 +353,7 @@ def data(self): def test_joint_state_shapes(self, data, property_name, expected_shape): """Test joint state properties have correct shapes.""" prop = getattr(data, property_name) - assert wp.to_torch(prop).shape == expected_shape, f"{property_name} has wrong shape" + assert prop.torch.shape == expected_shape, f"{property_name} has wrong shape" # -- Joint Property Shapes -- @pytest.mark.parametrize( @@ -385,7 +376,7 @@ def test_joint_state_shapes(self, data, property_name, expected_shape): def test_joint_property_shapes(self, data, property_name, expected_shape): """Test joint property shapes.""" prop = getattr(data, property_name) - assert wp.to_torch(prop).shape == expected_shape, f"{property_name} has wrong shape" + assert prop.torch.shape == expected_shape, f"{property_name} has wrong shape" # -- Root State Properties -- @pytest.mark.parametrize( @@ -411,7 +402,7 @@ def test_joint_property_shapes(self, data, property_name, expected_shape): def test_root_state_shapes(self, data, property_name, expected_shape): """Test root state properties have correct shapes.""" prop = getattr(data, property_name) - assert wp.to_torch(prop).shape == expected_shape, f"{property_name} has wrong shape" + assert prop.torch.shape == expected_shape, f"{property_name} has wrong shape" # -- Body State Properties -- @pytest.mark.parametrize( @@ -439,7 +430,7 @@ def test_root_state_shapes(self, data, property_name, expected_shape): def test_body_state_shapes(self, data, property_name, expected_shape): """Test body state properties have correct shapes.""" prop = getattr(data, property_name) - assert wp.to_torch(prop).shape == expected_shape, f"{property_name} has wrong shape" + assert prop.torch.shape == expected_shape, f"{property_name} has wrong shape" # -- Body Properties -- @pytest.mark.parametrize( @@ -453,7 +444,7 @@ def test_body_state_shapes(self, data, property_name, expected_shape): def test_body_property_shapes(self, data, property_name, expected_shape): """Test body property shapes.""" prop = getattr(data, property_name) - assert wp.to_torch(prop).shape == expected_shape, f"{property_name} has wrong shape" + assert prop.torch.shape == expected_shape, f"{property_name} has wrong shape" # -- Default State Properties -- @pytest.mark.parametrize( @@ -469,7 +460,7 @@ def test_body_property_shapes(self, data, property_name, expected_shape): def test_default_state_shapes(self, data, property_name, expected_shape): """Test default state properties have correct shapes.""" prop = getattr(data, property_name) - assert wp.to_torch(prop).shape == expected_shape, f"{property_name} has wrong shape" + assert prop.torch.shape == expected_shape, f"{property_name} has wrong shape" # -- Derived Properties -- @pytest.mark.parametrize( @@ -486,7 +477,7 @@ def test_default_state_shapes(self, data, property_name, expected_shape): def test_derived_property_shapes(self, data, property_name, expected_shape): """Test derived properties have correct shapes.""" prop = getattr(data, property_name) - assert wp.to_torch(prop).shape == expected_shape, f"{property_name} has wrong shape" + assert prop.torch.shape == expected_shape, f"{property_name} has wrong shape" # -- Tendon Properties -- @pytest.mark.parametrize( @@ -507,7 +498,7 @@ def test_derived_property_shapes(self, data, property_name, expected_shape): def test_tendon_property_shapes(self, data, property_name, expected_shape): """Test tendon properties have correct shapes.""" prop = getattr(data, property_name) - assert wp.to_torch(prop).shape == expected_shape, f"{property_name} has wrong shape" + assert prop.torch.shape == expected_shape, f"{property_name} has wrong shape" # -- Setter Tests -- @pytest.mark.parametrize( @@ -525,12 +516,11 @@ def test_tendon_property_shapes(self, data, property_name, expected_shape): ) def test_setters_update_properties(self, data, setter_name, property_name, shape): """Test that setters properly update the corresponding properties.""" - import warp as wp test_value = torch.randn(shape) setter = getattr(data, setter_name) setter(test_value) - result = wp.to_torch(getattr(data, property_name)) + result = getattr(data, property_name).torch assert torch.allclose(result, test_value), f"{setter_name} did not update {property_name}" def test_bulk_setter_with_multiple_properties(self, data): @@ -545,9 +535,9 @@ def test_bulk_setter_with_multiple_properties(self, data): root_link_pose_w=root_pose, ) - assert torch.allclose(wp.to_torch(data.joint_pos), joint_pos) - assert torch.allclose(wp.to_torch(data.joint_vel), joint_vel) - assert torch.allclose(wp.to_torch(data.root_link_pose_w), root_pose) + assert torch.allclose(data.joint_pos.torch, joint_pos) + assert torch.allclose(data.joint_vel.torch, joint_vel) + assert torch.allclose(data.root_link_pose_w.torch, root_pose) def test_bulk_setter_unknown_property_raises(self, data): """Test that set_mock_data raises on unknown property.""" @@ -601,7 +591,7 @@ def data(self): 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 wp.to_torch(prop).shape == expected_shape, f"{property_name} has wrong shape" + assert prop.torch.shape == expected_shape, f"{property_name} has wrong shape" # ============================================================================== @@ -660,7 +650,7 @@ def data(self): 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 wp.to_torch(prop).shape == expected_shape, f"{property_name} has wrong shape" + assert prop.torch.shape == expected_shape, f"{property_name} has wrong shape" # ============================================================================== @@ -714,7 +704,7 @@ def test_pva_pose_composition(self): data.set_pos_w(pos) data.set_quat_w(quat) - pose = wp.to_torch(data.pose_w) + pose = data.pose_w.torch assert torch.allclose(pose[:, :3], pos) assert torch.allclose(pose[:, 3:], quat) @@ -727,7 +717,7 @@ def test_articulation_root_state_composition(self): data.set_root_link_pose_w(pose) data.set_root_link_vel_w(vel) - state = wp.to_torch(data.root_link_state_w) + state = data.root_link_state_w.torch assert torch.allclose(state[:, :7], pose) assert torch.allclose(state[:, 7:], vel) @@ -740,7 +730,7 @@ def test_articulation_default_state_composition(self): data.set_default_root_pose(pose) data.set_default_root_vel(vel) - state = wp.to_torch(data.default_root_state) + state = data.default_root_state.torch assert torch.allclose(state[:, :7], pose) assert torch.allclose(state[:, 7:], vel) @@ -782,8 +772,8 @@ 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 wp.to_torch(alias_prop).shape == expected_shape, f"{alias} has wrong shape" - assert torch.allclose(wp.to_torch(alias_prop), wp.to_torch(source_prop)), f"{alias} should equal {source}" + assert alias_prop.torch.shape == expected_shape, f"{alias} has wrong shape" + assert torch.allclose(alias_prop.torch, source_prop.torch), f"{alias} should equal {source}" # -- Body state aliases (without _link_ or _com_ prefix) -- @pytest.mark.parametrize( @@ -805,8 +795,8 @@ 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 wp.to_torch(alias_prop).shape == expected_shape, f"{alias} has wrong shape" - assert torch.allclose(wp.to_torch(alias_prop), wp.to_torch(source_prop)), f"{alias} should equal {source}" + assert alias_prop.torch.shape == expected_shape, f"{alias} has wrong shape" + assert torch.allclose(alias_prop.torch, source_prop.torch), f"{alias} should equal {source}" # -- CoM in body frame -- @pytest.mark.parametrize( @@ -819,7 +809,7 @@ def test_body_aliases(self, data, alias, source, expected_shape): def test_com_body_frame_aliases(self, data, alias, expected_shape): """Test CoM in body frame aliases.""" prop = getattr(data, alias) - assert wp.to_torch(prop).shape == expected_shape, f"{alias} has wrong shape" + assert prop.torch.shape == expected_shape, f"{alias} has wrong shape" # -- Joint property aliases -- @pytest.mark.parametrize( @@ -834,14 +824,14 @@ 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 wp.to_torch(alias_prop).shape == expected_shape, f"{alias} has wrong shape" - assert torch.allclose(wp.to_torch(alias_prop), wp.to_torch(source_prop)), f"{alias} should equal {source}" + assert alias_prop.torch.shape == expected_shape, f"{alias} has wrong shape" + assert torch.allclose(alias_prop.torch, source_prop.torch), 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 wp.to_torch(data.fixed_tendon_limit).shape == (4, 2, 2) - assert torch.allclose(wp.to_torch(data.fixed_tendon_limit), wp.to_torch(data.fixed_tendon_pos_limits)) + assert data.fixed_tendon_limit.torch.shape == (4, 2, 2) + assert torch.allclose(data.fixed_tendon_limit.torch, data.fixed_tendon_pos_limits.torch) class TestRigidObjectConvenienceAliases: @@ -879,8 +869,8 @@ 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 wp.to_torch(alias_prop).shape == expected_shape, f"{alias} has wrong shape" - assert torch.allclose(wp.to_torch(alias_prop), wp.to_torch(source_prop)), f"{alias} should equal {source}" + assert alias_prop.torch.shape == expected_shape, f"{alias} has wrong shape" + assert torch.allclose(alias_prop.torch, source_prop.torch), f"{alias} should equal {source}" @pytest.mark.parametrize( "alias,expected_shape", @@ -892,7 +882,7 @@ def test_aliases(self, data, alias, source, expected_shape): def test_com_body_frame_aliases(self, data, alias, expected_shape): """Test CoM in body frame aliases.""" prop = getattr(data, alias) - assert wp.to_torch(prop).shape == expected_shape, f"{alias} has wrong shape" + assert prop.torch.shape == expected_shape, f"{alias} has wrong shape" class TestRigidObjectCollectionConvenienceAliases: @@ -909,5 +899,5 @@ def data(self): def test_body_state_w_alias(self, data): """Test body_state_w is alias for body_com_state_w.""" - assert wp.to_torch(data.body_state_w).shape == (4, 5, 13) - assert torch.allclose(wp.to_torch(data.body_state_w), wp.to_torch(data.body_com_state_w)) + assert data.body_state_w.torch.shape == (4, 5, 13) + assert torch.allclose(data.body_state_w.torch, data.body_com_state_w.torch) diff --git a/source/isaaclab/test/test_mock_interfaces/test_mock_sensors.py b/source/isaaclab/test/test_mock_interfaces/test_mock_sensors.py index daa03a79887b..949677f1c398 100644 --- a/source/isaaclab/test/test_mock_interfaces/test_mock_sensors.py +++ b/source/isaaclab/test/test_mock_interfaces/test_mock_sensors.py @@ -13,6 +13,7 @@ import pytest import torch +import warp as wp from isaaclab.test.mock_interfaces.sensors import ( MockContactSensor, @@ -50,40 +51,34 @@ def test_initialization(self, imu): def test_lazy_tensor_initialization(self, imu): """Test that unset properties return zero tensors with correct shapes.""" - import warp as wp - # IMU only has angular velocity and linear acceleration - ang_vel = wp.to_torch(imu.data.ang_vel_b) + ang_vel = imu.data.ang_vel_b.torch assert ang_vel.shape == (4, 3) assert torch.all(ang_vel == 0) - lin_acc = wp.to_torch(imu.data.lin_acc_b) + lin_acc = imu.data.lin_acc_b.torch assert lin_acc.shape == (4, 3) assert torch.all(lin_acc == 0) def test_set_mock_data(self, imu): """Test bulk data setter.""" - import warp as wp - ang_vel = torch.randn(4, 3) lin_acc = torch.randn(4, 3) imu.data.set_mock_data(ang_vel_b=ang_vel, lin_acc_b=lin_acc) - assert torch.allclose(wp.to_torch(imu.data.ang_vel_b), ang_vel) - assert torch.allclose(wp.to_torch(imu.data.lin_acc_b), lin_acc) + assert torch.allclose(imu.data.ang_vel_b.torch, ang_vel) + assert torch.allclose(imu.data.lin_acc_b.torch, lin_acc) def test_per_property_setter(self, imu): """Test individual property setters.""" - import warp as wp - lin_acc = torch.randn(4, 3) imu.data.set_lin_acc_b(lin_acc) - assert torch.allclose(wp.to_torch(imu.data.lin_acc_b), lin_acc) + assert torch.allclose(imu.data.lin_acc_b.torch, lin_acc) ang_vel = torch.randn(4, 3) imu.data.set_ang_vel_b(ang_vel) - assert torch.allclose(wp.to_torch(imu.data.ang_vel_b), ang_vel) + assert torch.allclose(imu.data.ang_vel_b.torch, ang_vel) # ============================================================================== @@ -107,15 +102,13 @@ def test_initialization(self, pva): def test_lazy_tensor_initialization(self, pva): """Test that unset properties return zero tensors with correct shapes.""" - import warp as wp - # Position - pos = wp.to_torch(pva.data.pos_w) + pos = pva.data.pos_w.torch assert pos.shape == (4, 3) assert torch.all(pos == 0) # Quaternion (should be identity in XYZW format: x=0, y=0, z=0, w=1) - quat = wp.to_torch(pva.data.quat_w) + quat = pva.data.quat_w.torch assert quat.shape == (4, 4) assert torch.all(quat[:, :3] == 0) # xyz components assert torch.all(quat[:, 3] == 1) # w component @@ -128,44 +121,36 @@ def test_lazy_tensor_initialization(self, pva): def test_projected_gravity_default(self, pva): """Test default gravity direction.""" - import warp as wp - - gravity = wp.to_torch(pva.data.projected_gravity_b) + gravity = pva.data.projected_gravity_b.torch 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, pva): """Test bulk data setter.""" - import warp as wp - lin_vel = torch.randn(4, 3) ang_vel = torch.randn(4, 3) pva.data.set_mock_data(lin_vel_b=lin_vel, ang_vel_b=ang_vel) - assert torch.allclose(wp.to_torch(pva.data.lin_vel_b), lin_vel) - assert torch.allclose(wp.to_torch(pva.data.ang_vel_b), ang_vel) + assert torch.allclose(pva.data.lin_vel_b.torch, lin_vel) + assert torch.allclose(pva.data.ang_vel_b.torch, ang_vel) def test_per_property_setter(self, pva): """Test individual property setters.""" - import warp as wp - lin_acc = torch.randn(4, 3) pva.data.set_lin_acc_b(lin_acc) - assert torch.allclose(wp.to_torch(pva.data.lin_acc_b), lin_acc) + assert torch.allclose(pva.data.lin_acc_b.torch, lin_acc) def test_pose_composition(self, pva): """Test that pose_w combines pos_w and quat_w correctly.""" - import warp as wp - pos = torch.randn(4, 3) quat = torch.tensor([[0, 0, 0, 1]] * 4, dtype=torch.float32) # XYZW format pva.data.set_pos_w(pos) pva.data.set_quat_w(quat) - pose = wp.to_torch(pva.data.pose_w) + pose = pva.data.pose_w.torch assert pose.shape == (4, 7) assert torch.allclose(pose[:, :3], pos) assert torch.allclose(pose[:, 3:], quat) @@ -227,7 +212,6 @@ def test_find_bodies(self, sensor): def test_compute_first_contact(self, sensor): """Test first contact computation.""" - import warp as wp # Set contact time to 0.5 for all bodies sensor.data.set_current_contact_time(torch.full((4, 4), 0.5)) @@ -242,7 +226,6 @@ def test_compute_first_contact(self, sensor): def test_compute_first_air(self, sensor): """Test first air computation.""" - import warp as wp sensor.data.set_current_air_time(torch.full((4, 4), 0.2)) diff --git a/source/isaaclab/test/test_scripts_torcharray_patterns.py b/source/isaaclab/test/test_scripts_torcharray_patterns.py new file mode 100644 index 000000000000..227bbda7be3c --- /dev/null +++ b/source/isaaclab/test/test_scripts_torcharray_patterns.py @@ -0,0 +1,133 @@ +# 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 + +"""Static scanner ensuring scripts/source do not regress the ProxyArray migration. + +Every public ``.data.`` on an asset or sensor now returns a +:class:`~isaaclab.utils.warp.ProxyArray` (or, for CameraData, a raw +``torch.Tensor``). Legacy conversion or tensor-method callsites on these +properties should migrate to explicit ``.torch`` or ``.warp`` access. These +tests regex-scan user-facing scripts/source files and flag regressions before +they reach users running tutorials or demos. +""" + +from __future__ import annotations + +import re +from pathlib import Path + +import pytest + +# Matches: +# wp.to_torch(.data.) +# where can be , [...], ., .[...], etc. +# Anchored so we don't match tool scripts that document the pattern as prose. +_WP_TO_TORCH_DOT_DATA = re.compile( + r"wp\.to_torch\(\s*" + r"[a-zA-Z_][a-zA-Z_0-9]*" # first name + r"(?:\[[^\]\[]*\]|\.[a-zA-Z_][a-zA-Z_0-9]*)*" # chain of [..] or .name + r"\.data\." # .data. + r"[a-zA-Z_][a-zA-Z_0-9]*" # field + r"\s*\)" +) + +# Matches: +# wp.to_torch(_data.) +# Catches the aliased pattern like ``object_data: RigidObjectData = ...`` followed +# by ``wp.to_torch(object_data.root_pos_w)``. +_WP_TO_TORCH_NAME_DATA = re.compile( + r"wp\.to_torch\(\s*" + r"[a-zA-Z_][a-zA-Z_0-9]*_data\." + r"[a-zA-Z_][a-zA-Z_0-9]*" + r"\s*\)" +) + +# Matches: +# .data..clone() +# .data.[...].clone() +# .data..assign(...) +# These are tensor/wp.array instance methods that ProxyArray intentionally does +# not forward. ``data.output[...]`` is camera data and remains torch-native. +_PROXYARRAY_DIRECT_METHOD_DOT_DATA = re.compile( + r"\.data\." + r"(?!_)" # ignore private backing buffers such as data._sim_bind_... + r"(?!output\b)" # camera output dict is torch-native + r"[a-zA-Z_][a-zA-Z_0-9]*" + r"(?:\[[^\]\[]*\])?" + r"\.(?:clone|assign)\s*\(" +) + +# scripts/tools/wrap_warp_to_torch.py is the migration utility and documents +# the old pattern inside strings. Exclude it from the scan. +_EXCLUDE = {"tools/wrap_warp_to_torch.py"} +_EXCLUDE_PREFIXES = ("source/isaaclab_contrib/",) + + +def _repo_root() -> Path: + # this test lives at source/isaaclab/test/test_scripts_torcharray_patterns.py + # parents[0]=test, [1]=isaaclab, [2]=source, [3]=repo root + return Path(__file__).resolve().parents[3] + + +def _scripts_files() -> list[Path]: + scripts = _repo_root() / "scripts" + return sorted(p for p in scripts.rglob("*.py") if "__pycache__" not in p.parts) + + +def _source_and_scripts_files() -> list[Path]: + roots = [_repo_root() / "scripts", _repo_root() / "source"] + return sorted(p for root in roots for p in root.rglob("*.py") if "__pycache__" not in p.parts) + + +@pytest.mark.parametrize("path", _scripts_files(), ids=lambda p: str(p.relative_to(_repo_root()))) +def test_no_wp_to_torch_on_torcharray_data(path: Path) -> None: + """No ``wp.to_torch(.data.)`` / ``wp.to_torch(_data.)`` in scripts/. + + Post-migration, ``.data.`` returns a ``ProxyArray`` + (or ``torch.Tensor`` for CameraData). The temporary ``wp.to_torch`` + shim is deprecated, so use the ``.torch`` accessor instead (or omit + the wrap entirely for torch-native fields). + """ + rel = path.relative_to(_repo_root()).as_posix() + if any(rel.endswith(suffix) for suffix in _EXCLUDE): + pytest.skip(f"{rel} is excluded from the ProxyArray hygiene scan") + if rel.startswith(_EXCLUDE_PREFIXES): + pytest.skip(f"{rel} is outside the ProxyArray migration scan scope") + + text = path.read_text(encoding="utf-8") + + offenders: list[str] = [] + for i, line in enumerate(text.splitlines(), start=1): + if _WP_TO_TORCH_DOT_DATA.search(line) or _WP_TO_TORCH_NAME_DATA.search(line): + offenders.append(f"{rel}:{i}: {line.rstrip()}") + + if offenders: + pytest.fail( + "Found wp.to_torch(...) calls on a migrated ProxyArray data accessor. " + "Use .torch instead of wp.to_torch(...) (see isaaclab 4.6.15 CHANGELOG).\n" + "\n".join(offenders) + ) + + +@pytest.mark.parametrize("path", _source_and_scripts_files(), ids=lambda p: str(p.relative_to(_repo_root()))) +def test_no_direct_proxyarray_data_methods(path: Path) -> None: + """No direct tensor/wp.array methods on migrated ``.data.`` accessors.""" + rel = path.relative_to(_repo_root()).as_posix() + if any(rel.endswith(suffix) for suffix in _EXCLUDE): + pytest.skip(f"{rel} is excluded from the ProxyArray hygiene scan") + if rel.startswith(_EXCLUDE_PREFIXES): + pytest.skip(f"{rel} is outside the ProxyArray migration scan scope") + + text = path.read_text(encoding="utf-8") + + offenders: list[str] = [] + for i, line in enumerate(text.splitlines(), start=1): + if _PROXYARRAY_DIRECT_METHOD_DOT_DATA.search(line): + offenders.append(f"{rel}:{i}: {line.rstrip()}") + + if offenders: + pytest.fail( + "Found direct tensor/wp.array methods on migrated ProxyArray data accessors. " + "Use .torch.clone() for tensor copies or .warp.assign(...) for warp writes.\n" + "\n".join(offenders) + ) diff --git a/source/isaaclab/test/utils/test_wrench_composer.py b/source/isaaclab/test/utils/test_wrench_composer.py index 37f6d5959aac..cffa927c802b 100644 --- a/source/isaaclab/test/utils/test_wrench_composer.py +++ b/source/isaaclab/test/utils/test_wrench_composer.py @@ -137,7 +137,7 @@ def test_wrench_composer_add_force(device: str, num_envs: int, num_bodies: int): # Compose to body frame before checking output wrench_composer.compose_to_body_frame() # Get composed force from wrench composer - composed_force_np = wrench_composer.out_force_b.numpy() + composed_force_np = wrench_composer.out_force_b.warp.numpy() assert np.allclose(composed_force_np, hand_calculated_composed_force_np, atol=1, rtol=1e-7) @@ -176,7 +176,7 @@ def test_wrench_composer_add_torque(device: str, num_envs: int, num_bodies: int) # Compose to body frame before checking output wrench_composer.compose_to_body_frame() # Get composed torque from wrench composer - composed_torque_np = wrench_composer.out_torque_b.numpy() + composed_torque_np = wrench_composer.out_torque_b.warp.numpy() assert np.allclose(composed_torque_np, hand_calculated_composed_torque_np, atol=1, rtol=1e-7) @@ -232,10 +232,10 @@ def test_add_forces_at_positions(device: str, num_envs: int, num_bodies: int): # Compose to body frame before checking output wrench_composer.compose_to_body_frame() # Get composed force from wrench composer - composed_force_np = wrench_composer.out_force_b.numpy() + composed_force_np = wrench_composer.out_force_b.warp.numpy() assert np.allclose(composed_force_np, hand_calculated_composed_force_np, atol=1, rtol=1e-7) # Get composed torque from wrench composer - composed_torque_np = wrench_composer.out_torque_b.numpy() + composed_torque_np = wrench_composer.out_torque_b.warp.numpy() assert np.allclose(composed_torque_np, hand_calculated_composed_torque_np, atol=1, rtol=1e-7) @@ -281,7 +281,7 @@ def test_add_torques_at_position(device: str, num_envs: int, num_bodies: int): # Compose to body frame before checking output wrench_composer.compose_to_body_frame() # Get composed torque from wrench composer - composed_torque_np = wrench_composer.out_torque_b.numpy() + composed_torque_np = wrench_composer.out_torque_b.warp.numpy() assert np.allclose(composed_torque_np, hand_calculated_composed_torque_np, atol=1, rtol=1e-7) @@ -341,10 +341,10 @@ def test_add_forces_and_torques_at_position(device: str, num_envs: int, num_bodi # Compose to body frame before checking output wrench_composer.compose_to_body_frame() # Get composed force from wrench composer - composed_force_np = wrench_composer.out_force_b.numpy() + composed_force_np = wrench_composer.out_force_b.warp.numpy() assert np.allclose(composed_force_np, hand_calculated_composed_force_np, atol=1, rtol=1e-7) # Get composed torque from wrench composer - composed_torque_np = wrench_composer.out_torque_b.numpy() + composed_torque_np = wrench_composer.out_torque_b.warp.numpy() assert np.allclose(composed_torque_np, hand_calculated_composed_torque_np, atol=1, rtol=1e-7) @@ -388,8 +388,8 @@ def test_wrench_composer_reset(device: str, num_envs: int, num_bodies: int): assert np.allclose(wrench_composer.global_force_at_com_w.numpy(), zeros, atol=1, rtol=1e-7) assert np.allclose(wrench_composer.local_force_b.numpy(), zeros, atol=1, rtol=1e-7) assert np.allclose(wrench_composer.local_torque_b.numpy(), zeros, atol=1, rtol=1e-7) - assert np.allclose(wrench_composer.out_force_b.numpy(), zeros, atol=1, rtol=1e-7) - assert np.allclose(wrench_composer.out_torque_b.numpy(), zeros, atol=1, rtol=1e-7) + assert np.allclose(wrench_composer.out_force_b.warp.numpy(), zeros, atol=1, rtol=1e-7) + assert np.allclose(wrench_composer.out_torque_b.warp.numpy(), zeros, atol=1, rtol=1e-7) # ============================================================================ @@ -433,7 +433,7 @@ def test_global_forces_with_rotation(device: str, num_envs: int, num_bodies: int wrench_composer.compose_to_body_frame() # Verify - composed_force_np = wrench_composer.out_force_b.numpy() + composed_force_np = wrench_composer.out_force_b.warp.numpy() assert np.allclose(composed_force_np, expected_forces_local, atol=1e-4, rtol=1e-5), ( f"Global force rotation failed.\nExpected:\n{expected_forces_local}\nGot:\n{composed_force_np}" ) @@ -475,7 +475,7 @@ def test_global_torques_with_rotation(device: str, num_envs: int, num_bodies: in wrench_composer.compose_to_body_frame() # Verify - composed_torque_np = wrench_composer.out_torque_b.numpy() + composed_torque_np = wrench_composer.out_torque_b.warp.numpy() assert np.allclose(composed_torque_np, expected_torques_local, atol=1e-4, rtol=1e-5), ( f"Global torque rotation failed.\nExpected:\n{expected_torques_local}\nGot:\n{composed_torque_np}" ) @@ -533,13 +533,13 @@ def test_global_forces_at_global_position(device: str, num_envs: int, num_bodies wrench_composer.compose_to_body_frame() # Verify forces - composed_force_np = wrench_composer.out_force_b.numpy() + composed_force_np = wrench_composer.out_force_b.warp.numpy() assert np.allclose(composed_force_np, expected_forces_local, atol=1e-3, rtol=1e-4), ( f"Global force at position failed.\nExpected forces:\n{expected_forces_local}\nGot:\n{composed_force_np}" ) # Verify torques - composed_torque_np = wrench_composer.out_torque_b.numpy() + composed_torque_np = wrench_composer.out_torque_b.warp.numpy() assert np.allclose(composed_torque_np, expected_torques_local, atol=1e-3, rtol=1e-4), ( f"Global force at position failed.\nExpected torques:\n{expected_torques_local}\nGot:\n{composed_torque_np}" ) @@ -576,13 +576,13 @@ def test_local_vs_global_identity_quaternion(device: str): # Results should be identical assert np.allclose( - wrench_composer_local.out_force_b.numpy(), - wrench_composer_global.out_force_b.numpy(), + wrench_composer_local.out_force_b.warp.numpy(), + wrench_composer_global.out_force_b.warp.numpy(), atol=1e-6, ) assert np.allclose( - wrench_composer_local.out_torque_b.numpy(), - wrench_composer_global.out_torque_b.numpy(), + wrench_composer_local.out_torque_b.warp.numpy(), + wrench_composer_global.out_torque_b.warp.numpy(), atol=1e-6, ) @@ -614,7 +614,7 @@ def test_90_degree_rotation_global_force(device: str): # Compose to body frame before checking output wrench_composer.compose_to_body_frame() - composed_force_np = wrench_composer.out_force_b.numpy() + composed_force_np = wrench_composer.out_force_b.warp.numpy() assert np.allclose(composed_force_np, expected_force_local, atol=1e-5), ( f"90-degree rotation test failed.\nExpected:\n{expected_force_local}\nGot:\n{composed_force_np}" ) @@ -659,7 +659,7 @@ def test_composition_mixed_local_and_global(device: str): # Compose to body frame before checking output wrench_composer.compose_to_body_frame() - composed_force_np = wrench_composer.out_force_b.numpy() + composed_force_np = wrench_composer.out_force_b.warp.numpy() assert np.allclose(composed_force_np, expected_total, atol=1e-4, rtol=1e-5), ( f"Mixed local/global composition failed.\nExpected:\n{expected_total}\nGot:\n{composed_force_np}" ) @@ -703,8 +703,8 @@ def test_local_forces_at_local_position(device: str, num_envs: int, num_bodies: wrench_composer.compose_to_body_frame() # Verify - composed_force_np = wrench_composer.out_force_b.numpy() - composed_torque_np = wrench_composer.out_torque_b.numpy() + composed_force_np = wrench_composer.out_force_b.warp.numpy() + composed_torque_np = wrench_composer.out_torque_b.warp.numpy() assert np.allclose(composed_force_np, expected_forces, atol=1e-4, rtol=1e-5) assert np.allclose(composed_torque_np, expected_torques, atol=1e-4, rtol=1e-5) @@ -746,8 +746,8 @@ def test_global_force_at_link_origin_no_torque(device: str): # Compose to body frame before checking output wrench_composer.compose_to_body_frame() - composed_force_np = wrench_composer.out_force_b.numpy() - composed_torque_np = wrench_composer.out_torque_b.numpy() + composed_force_np = wrench_composer.out_force_b.warp.numpy() + composed_torque_np = wrench_composer.out_torque_b.warp.numpy() assert np.allclose(composed_force_np, expected_forces, atol=1e-4, rtol=1e-5) assert np.allclose(composed_torque_np, expected_torques, atol=1e-4, rtol=1e-5) @@ -816,12 +816,12 @@ def test_add_raw_buffers_from(device: str, num_envs: int, num_bodies: int): composer_a.compose_to_body_frame() composer_ref.compose_to_body_frame() - assert np.allclose(composer_a.out_force_b.numpy(), composer_ref.out_force_b.numpy(), atol=1e-4, rtol=1e-5), ( - "add_raw_buffers_from force mismatch vs direct accumulation" - ) - assert np.allclose(composer_a.out_torque_b.numpy(), composer_ref.out_torque_b.numpy(), atol=1e-4, rtol=1e-5), ( - "add_raw_buffers_from torque mismatch vs direct accumulation" - ) + assert np.allclose( + composer_a.out_force_b.warp.numpy(), composer_ref.out_force_b.warp.numpy(), atol=1e-4, rtol=1e-5 + ), "add_raw_buffers_from force mismatch vs direct accumulation" + assert np.allclose( + composer_a.out_torque_b.warp.numpy(), composer_ref.out_torque_b.warp.numpy(), atol=1e-4, rtol=1e-5 + ), "add_raw_buffers_from torque mismatch vs direct accumulation" @pytest.mark.parametrize("device", ["cuda:0", "cpu"]) @@ -905,11 +905,11 @@ def test_add_forces_mask(device: str, num_envs: int, num_bodies: int): composer_idx.compose_to_body_frame() composer_mask.compose_to_body_frame() - assert np.allclose(composer_idx.out_force_b.numpy(), composer_mask.out_force_b.numpy(), atol=1e-4, rtol=1e-5), ( - f"Mask vs index force mismatch (envs={num_envs}, bodies={num_bodies})" - ) assert np.allclose( - composer_idx.out_torque_b.numpy(), composer_mask.out_torque_b.numpy(), atol=1e-4, rtol=1e-5 + composer_idx.out_force_b.warp.numpy(), composer_mask.out_force_b.warp.numpy(), atol=1e-4, rtol=1e-5 + ), f"Mask vs index force mismatch (envs={num_envs}, bodies={num_bodies})" + assert np.allclose( + composer_idx.out_torque_b.warp.numpy(), composer_mask.out_torque_b.warp.numpy(), atol=1e-4, rtol=1e-5 ), f"Mask vs index torque mismatch (envs={num_envs}, bodies={num_bodies})" @@ -955,12 +955,12 @@ def test_add_forces_mask_global(device: str, num_envs: int, num_bodies: int): composer_idx.compose_to_body_frame() composer_mask.compose_to_body_frame() - assert np.allclose(composer_idx.out_force_b.numpy(), composer_mask.out_force_b.numpy(), atol=1e-4, rtol=1e-5), ( - "Mask vs index global force mismatch" - ) - assert np.allclose(composer_idx.out_torque_b.numpy(), composer_mask.out_torque_b.numpy(), atol=1e-4, rtol=1e-5), ( - "Mask vs index global torque mismatch" - ) + assert np.allclose( + composer_idx.out_force_b.warp.numpy(), composer_mask.out_force_b.warp.numpy(), atol=1e-4, rtol=1e-5 + ), "Mask vs index global force mismatch" + assert np.allclose( + composer_idx.out_torque_b.warp.numpy(), composer_mask.out_torque_b.warp.numpy(), atol=1e-4, rtol=1e-5 + ), "Mask vs index global torque mismatch" # ============================================================================ @@ -992,7 +992,7 @@ def test_set_forces_overwrites_previous_add(device: str): composer.compose_to_body_frame() # Output should match forces_b only (forces_a should be gone) - assert np.allclose(composer.out_force_b.numpy(), forces_b_np, atol=1e-4, rtol=1e-5), ( + assert np.allclose(composer.out_force_b.warp.numpy(), forces_b_np, atol=1e-4, rtol=1e-5), ( "set_forces did not clear previous add" ) @@ -1160,7 +1160,7 @@ def test_composed_force_emits_deprecation_warning(device: str): result = composer.composed_force # Should return the same data as out_force_b - assert np.allclose(result.numpy(), composer.out_force_b.numpy(), atol=1e-7) + assert np.allclose(result.warp.numpy(), composer.out_force_b.warp.numpy(), atol=1e-7) @pytest.mark.parametrize("device", ["cuda:0", "cpu"]) @@ -1179,7 +1179,7 @@ def test_composed_torque_emits_deprecation_warning(device: str): with pytest.warns(DeprecationWarning, match="composed_torque.*is deprecated"): result = composer.composed_torque - assert np.allclose(result.numpy(), composer.out_torque_b.numpy(), atol=1e-7) + assert np.allclose(result.warp.numpy(), composer.out_torque_b.warp.numpy(), atol=1e-7) @pytest.mark.parametrize("device", ["cuda:0", "cpu"]) @@ -1199,7 +1199,7 @@ def test_deprecated_add_forces_and_torques_emits_warning(device: str): ) composer.compose_to_body_frame() - assert np.allclose(composer.out_force_b.numpy(), forces_np, atol=1e-4, rtol=1e-5) + assert np.allclose(composer.out_force_b.warp.numpy(), forces_np, atol=1e-4, rtol=1e-5) # ============================================================================ @@ -1231,7 +1231,7 @@ def test_set_forces_mask_overwrites_previous_add(device: str): composer.compose_to_body_frame() # Output should match forces_b only (forces_a should be gone) - assert np.allclose(composer.out_force_b.numpy(), forces_b_np, atol=1e-4, rtol=1e-5), ( + assert np.allclose(composer.out_force_b.warp.numpy(), forces_b_np, atol=1e-4, rtol=1e-5), ( "set_forces_and_torques_mask did not clear previous add" ) @@ -1343,12 +1343,12 @@ def test_set_forces_mask_matches_set_forces_index(device: str): composer_idx.compose_to_body_frame() composer_mask.compose_to_body_frame() - assert np.allclose(composer_idx.out_force_b.numpy(), composer_mask.out_force_b.numpy(), atol=1e-4, rtol=1e-5), ( - "set mask vs index force mismatch" - ) - assert np.allclose(composer_idx.out_torque_b.numpy(), composer_mask.out_torque_b.numpy(), atol=1e-4, rtol=1e-5), ( - "set mask vs index torque mismatch" - ) + assert np.allclose( + composer_idx.out_force_b.warp.numpy(), composer_mask.out_force_b.warp.numpy(), atol=1e-4, rtol=1e-5 + ), "set mask vs index force mismatch" + assert np.allclose( + composer_idx.out_torque_b.warp.numpy(), composer_mask.out_torque_b.warp.numpy(), atol=1e-4, rtol=1e-5 + ), "set mask vs index torque mismatch" # ============================================================================ @@ -1376,7 +1376,7 @@ def test_out_force_b_triggers_lazy_composition(device: str): # Do NOT call compose_to_body_frame -- rely on lazy composition expected_forces_local = quat_rotate_inv_np(link_quat_np, forces_global_np) - composed_force_np = composer.out_force_b.numpy() + composed_force_np = composer.out_force_b.warp.numpy() assert np.allclose(composed_force_np, expected_forces_local, atol=1e-4, rtol=1e-5), ( "Lazy composition via out_force_b failed" @@ -1403,7 +1403,7 @@ def test_out_torque_b_triggers_lazy_composition(device: str): # Do NOT call compose_to_body_frame -- rely on lazy composition expected_torques_local = quat_rotate_inv_np(link_quat_np, torques_global_np) - composed_torque_np = composer.out_torque_b.numpy() + composed_torque_np = composer.out_torque_b.warp.numpy() assert np.allclose(composed_torque_np, expected_torques_local, atol=1e-4, rtol=1e-5), ( "Lazy composition via out_torque_b failed" @@ -1444,7 +1444,7 @@ def test_lazy_composition_tracks_dirty_flag(device: str): # Verify accumulated result (2x forces) expected = 2.0 * forces_np - assert np.allclose(composer.out_force_b.numpy(), expected, atol=1e-4, rtol=1e-5) + assert np.allclose(composer.out_force_b.warp.numpy(), expected, atol=1e-4, rtol=1e-5) @pytest.mark.parametrize("device", ["cuda:0", "cpu"]) @@ -1481,13 +1481,13 @@ def test_compose_is_idempotent(device: str): # First compose composer.compose_to_body_frame() - force_first = composer.out_force_b.numpy().copy() - torque_first = composer.out_torque_b.numpy().copy() + force_first = composer.out_force_b.warp.numpy().copy() + torque_first = composer.out_torque_b.warp.numpy().copy() # Second compose (no writes in between) composer.compose_to_body_frame() - force_second = composer.out_force_b.numpy() - torque_second = composer.out_torque_b.numpy() + force_second = composer.out_force_b.warp.numpy() + torque_second = composer.out_torque_b.warp.numpy() np.testing.assert_array_equal(force_first, force_second) np.testing.assert_array_equal(torque_first, torque_second) @@ -1544,12 +1544,13 @@ def test_global_force_with_com_offset(device: str): expected_torque = np.zeros((num_envs, num_bodies, 3), dtype=np.float32) expected_torque[..., 1] = 10.0 - assert np.allclose(composer.out_torque_b.numpy(), expected_torque, atol=1e-4, rtol=1e-5), ( - f"CoM offset torque correction failed.\nExpected:\n{expected_torque}\nGot:\n{composer.out_torque_b.numpy()}" + actual_torque = composer.out_torque_b.warp.numpy() + assert np.allclose(actual_torque, expected_torque, atol=1e-4, rtol=1e-5), ( + f"CoM offset torque correction failed.\nExpected:\n{expected_torque}\nGot:\n{actual_torque}" ) # Force should be unchanged (identity rotation) - assert np.allclose(composer.out_force_b.numpy(), forces_np, atol=1e-4, rtol=1e-5) + assert np.allclose(composer.out_force_b.warp.numpy(), forces_np, atol=1e-4, rtol=1e-5) @pytest.mark.parametrize("device", ["cuda:0", "cpu"]) @@ -1593,7 +1594,7 @@ def test_global_force_at_com_no_torque_with_com_offset(device: str): # Torque = cross(com, F) - cross(com, F) = 0 expected_torque = np.zeros((num_envs, num_bodies, 3), dtype=np.float32) - assert np.allclose(composer.out_torque_b.numpy(), expected_torque, atol=1e-4, rtol=1e-5), ( + assert np.allclose(composer.out_torque_b.warp.numpy(), expected_torque, atol=1e-4, rtol=1e-5), ( "Force at CoM should produce zero torque regardless of CoM offset" ) @@ -1642,12 +1643,12 @@ def test_com_offset_with_rotation(device: str): expected_torque_b = quat_rotate_inv_np(link_quat_np, torque_w) expected_force_b = quat_rotate_inv_np(link_quat_np, forces_np) - assert np.allclose(composer.out_force_b.numpy(), expected_force_b, atol=1e-3, rtol=1e-4), ( + assert np.allclose(composer.out_force_b.warp.numpy(), expected_force_b, atol=1e-3, rtol=1e-4), ( "Force mismatch with CoM offset + rotation" ) - assert np.allclose(composer.out_torque_b.numpy(), expected_torque_b, atol=1e-3, rtol=1e-4), ( + assert np.allclose(composer.out_torque_b.warp.numpy(), expected_torque_b, atol=1e-3, rtol=1e-4), ( f"Torque mismatch with CoM offset + rotation.\n" - f"Expected:\n{expected_torque_b}\nGot:\n{composer.out_torque_b.numpy()}" + f"Expected:\n{expected_torque_b}\nGot:\n{composer.out_torque_b.warp.numpy()}" ) @@ -1673,7 +1674,7 @@ def test_deprecated_set_forces_and_torques_emits_warning(device: str): ) composer.compose_to_body_frame() - assert np.allclose(composer.out_force_b.numpy(), forces_np, atol=1e-4, rtol=1e-5) + assert np.allclose(composer.out_force_b.warp.numpy(), forces_np, atol=1e-4, rtol=1e-5) @pytest.mark.parametrize("device", ["cuda:0", "cpu"]) @@ -1699,6 +1700,6 @@ def test_deprecated_set_forces_and_torques_clears_previous(device: str): ) composer.compose_to_body_frame() - assert np.allclose(composer.out_force_b.numpy(), forces_b_np, atol=1e-4, rtol=1e-5), ( + assert np.allclose(composer.out_force_b.warp.numpy(), forces_b_np, atol=1e-4, rtol=1e-5), ( "Deprecated set_forces_and_torques did not replace previous values" ) diff --git a/source/isaaclab/test/utils/test_wrench_composer_integration.py b/source/isaaclab/test/utils/test_wrench_composer_integration.py index f71690d96541..bb195e856030 100644 --- a/source/isaaclab/test/utils/test_wrench_composer_integration.py +++ b/source/isaaclab/test/utils/test_wrench_composer_integration.py @@ -72,7 +72,7 @@ def test_global_force_invariant_under_rotation(device): body_ids, _ = cube_object.find_bodies(".*") mass = float(wp.to_torch(cube_object.root_view.get_masses())[0]) - com = wp.to_torch(cube_object.data.body_com_pos_w).clone() + com = cube_object.data.body_com_pos_w.torch.clone() # Apply permanent global force along +X at CoM forces = torch.zeros(1, len(body_ids), 3, device=device) @@ -93,10 +93,10 @@ def test_global_force_invariant_under_rotation(device): sim.step() cube_object.update(sim.cfg.dt) - vel_after_phase1 = wp.to_torch(cube_object.data.root_lin_vel_w)[0].clone() + vel_after_phase1 = cube_object.data.root_lin_vel_w.torch[0].clone() # Rotate body 180deg about Z (quat wxyz = [0, 0, 0, 1]) while keeping velocity - root_pose = wp.to_torch(cube_object.data.root_state_w)[0, :7].clone().unsqueeze(0) + root_pose = cube_object.data.root_state_w.torch[0, :7].clone().unsqueeze(0) root_pose[0, 3:7] = torch.tensor([0.0, 0.0, 1.0, 0.0], device=device) # 180deg about Z (xyzw) cube_object.write_root_pose_to_sim(root_pose) @@ -106,7 +106,7 @@ def test_global_force_invariant_under_rotation(device): sim.step() cube_object.update(sim.cfg.dt) - vel_after_phase2 = wp.to_torch(cube_object.data.root_lin_vel_w)[0].clone() + vel_after_phase2 = cube_object.data.root_lin_vel_w.torch[0].clone() # Acceleration should be same in both phases: delta_v_phase2 ≈ delta_v_phase1 delta_v_phase1 = vel_after_phase1[0].item() # vx after phase 1 @@ -165,11 +165,11 @@ def test_local_force_follows_rotation(device): sim.step() cube_object.update(sim.cfg.dt) - vel_after_phase1 = wp.to_torch(cube_object.data.root_lin_vel_w)[0].clone() + vel_after_phase1 = cube_object.data.root_lin_vel_w.torch[0].clone() assert vel_after_phase1[0].item() > 1.0, "Object should be moving in +X" # Rotate body 180deg about Z while keeping velocity - root_pose = wp.to_torch(cube_object.data.root_state_w)[0, :7].clone().unsqueeze(0) + root_pose = cube_object.data.root_state_w.torch[0, :7].clone().unsqueeze(0) root_pose[0, 3:7] = torch.tensor([0.0, 0.0, 1.0, 0.0], device=device) # 180deg about Z (xyzw) cube_object.write_root_pose_to_sim(root_pose) @@ -179,7 +179,7 @@ def test_local_force_follows_rotation(device): sim.step() cube_object.update(sim.cfg.dt) - vel_after_phase2 = wp.to_torch(cube_object.data.root_lin_vel_w)[0].clone() + vel_after_phase2 = cube_object.data.root_lin_vel_w.torch[0].clone() # Velocity should be approximately zero: decelerated by the same amount as it accelerated torch.testing.assert_close( @@ -213,7 +213,7 @@ def test_global_force_at_offset_generates_torque(device): torques = torch.zeros(1, len(body_ids), 3, device=device) # Position offset: CoM position + 1m in Y (global frame) - com_pos = wp.to_torch(cube_object.data.body_com_pos_w)[:, body_ids, :3].clone() + com_pos = cube_object.data.body_com_pos_w.torch[:, body_ids, :3].clone() positions = com_pos.clone() positions[..., 1] += 1.0 # +1m Y offset @@ -231,8 +231,8 @@ def test_global_force_at_offset_generates_torque(device): sim.step() cube_object.update(sim.cfg.dt) - lin_vel = wp.to_torch(cube_object.data.root_lin_vel_w)[0] - ang_vel = wp.to_torch(cube_object.data.root_ang_vel_w)[0] + lin_vel = cube_object.data.root_lin_vel_w.torch[0] + ang_vel = cube_object.data.root_ang_vel_w.torch[0] # Linear velocity in +X should be positive assert lin_vel[0].item() > 0.1, f"Expected positive X velocity, got {lin_vel[0].item()}" @@ -275,11 +275,11 @@ def test_global_torque_invariant_under_rotation(device): sim.step() cube_object.update(sim.cfg.dt) - omega_z_after_phase1 = wp.to_torch(cube_object.data.root_ang_vel_w)[0, 2].clone().item() + omega_z_after_phase1 = cube_object.data.root_ang_vel_w.torch[0, 2].clone().item() # Rotate body 90deg about X and zero out velocities so phase 2 starts from rest # (avoids gyroscopic cross-coupling at high omega) - root_pose = wp.to_torch(cube_object.data.root_state_w)[0, :7].clone().unsqueeze(0) + root_pose = cube_object.data.root_state_w.torch[0, :7].clone().unsqueeze(0) root_pose[0, 3:7] = torch.tensor([0.7071, 0.0, 0.0, 0.7071], device=device) # 90deg about X (xyzw) cube_object.write_root_pose_to_sim(root_pose) cube_object.write_root_velocity_to_sim(torch.zeros(1, 6, device=device)) @@ -290,7 +290,7 @@ def test_global_torque_invariant_under_rotation(device): sim.step() cube_object.update(sim.cfg.dt) - omega_z_after_phase2 = wp.to_torch(cube_object.data.root_ang_vel_w)[0, 2].clone().item() + omega_z_after_phase2 = cube_object.data.root_ang_vel_w.torch[0, 2].clone().item() # Both phases start from rest — angular acceleration about Z should be the same torch.testing.assert_close( @@ -324,7 +324,7 @@ def test_global_force_torque_after_translation(device): body_ids, _ = cube_object.find_bodies(".*") # Phase 1 setup: Move cube to (1, 0, 1) and apply force at (1, 0, 1) - root_state = wp.to_torch(cube_object.data.root_state_w).clone() + root_state = cube_object.data.root_state_w.torch.clone() root_state[0, 0] = 1.0 # x = 1 root_state[0, 1] = 0.0 # y = 0 root_state[0, 2] = 1.0 # z = 1 @@ -337,7 +337,7 @@ def test_global_force_torque_after_translation(device): cube_object.update(sim.cfg.dt) # Get current CoM position for the force application point - com_pos = wp.to_torch(cube_object.data.body_com_pos_w)[:, body_ids, :3].clone() + com_pos = cube_object.data.body_com_pos_w.torch[:, body_ids, :3].clone() forces = torch.zeros(1, len(body_ids), 3, device=device) forces[..., 1] = FORCE_MAGNITUDE # +Y force @@ -357,8 +357,8 @@ def test_global_force_torque_after_translation(device): sim.step() cube_object.update(sim.cfg.dt) - ang_vel_phase1 = wp.to_torch(cube_object.data.root_ang_vel_w)[0].clone() - lin_vel_phase1 = wp.to_torch(cube_object.data.root_lin_vel_w)[0].clone() + ang_vel_phase1 = cube_object.data.root_ang_vel_w.torch[0].clone() + lin_vel_phase1 = cube_object.data.root_lin_vel_w.torch[0].clone() # Should have linear velocity in +Y assert lin_vel_phase1[1].item() > 0.1, f"Expected positive Y velocity, got {lin_vel_phase1[1].item()}" @@ -369,7 +369,7 @@ def test_global_force_torque_after_translation(device): ) # Phase 2: Teleport cube to origin, zero velocity, don't re-apply force - root_state2 = wp.to_torch(cube_object.data.root_state_w).clone() + root_state2 = cube_object.data.root_state_w.torch.clone() root_state2[0, 0] = 0.0 # x = 0 root_state2[0, 1] = 0.0 root_state2[0, 2] = 1.0 # z = 1 @@ -390,7 +390,7 @@ def test_global_force_torque_after_translation(device): sim.step() cube_object.update(sim.cfg.dt) - ang_vel_phase2 = wp.to_torch(cube_object.data.root_ang_vel_w)[0].clone() + ang_vel_phase2 = cube_object.data.root_ang_vel_w.torch[0].clone() # The X component of position changed from ~1 to ~0, so torque about Z changes. # stored_torque_z = com_x * Fy = ~1 * 10 = ~10 @@ -422,7 +422,7 @@ def test_global_force_torque_reverses_on_opposite_side(device): body_ids, _ = cube_object.find_bodies(".*") # Move cube to (-1, 0, 1) - root_state = wp.to_torch(cube_object.data.root_state_w).clone() + root_state = cube_object.data.root_state_w.torch.clone() root_state[0, 0] = -1.0 root_state[0, 1] = 0.0 root_state[0, 2] = 1.0 @@ -453,11 +453,11 @@ def test_global_force_torque_reverses_on_opposite_side(device): sim.step() cube_object.update(sim.cfg.dt) - omega_z_phase1 = wp.to_torch(cube_object.data.root_ang_vel_w)[0, 2].item() + omega_z_phase1 = cube_object.data.root_ang_vel_w.torch[0, 2].item() assert omega_z_phase1 > 0.1, f"Phase 1: expected positive omega_z, got {omega_z_phase1}" # Phase 2: Teleport cube to (+1, 0, 1), zero velocity - root_state2 = wp.to_torch(cube_object.data.root_state_w).clone() + root_state2 = cube_object.data.root_state_w.torch.clone() root_state2[0, 0] = 1.0 root_state2[0, 1] = 0.0 root_state2[0, 2] = 1.0 @@ -473,7 +473,7 @@ def test_global_force_torque_reverses_on_opposite_side(device): sim.step() cube_object.update(sim.cfg.dt) - omega_z_phase2 = wp.to_torch(cube_object.data.root_ang_vel_w)[0, 2].item() + omega_z_phase2 = cube_object.data.root_ang_vel_w.torch[0, 2].item() assert omega_z_phase2 < -0.1, f"Phase 2: expected negative omega_z, got {omega_z_phase2}" @@ -493,7 +493,7 @@ def test_global_force_no_position_no_torque(device): body_ids, _ = cube_object.find_bodies(".*") # Move cube to (2, 0, 1) - root_state = wp.to_torch(cube_object.data.root_state_w).clone() + root_state = cube_object.data.root_state_w.torch.clone() root_state[0, 0] = 2.0 root_state[0, 1] = 0.0 root_state[0, 2] = 1.0 @@ -521,12 +521,12 @@ def test_global_force_no_position_no_torque(device): sim.step() cube_object.update(sim.cfg.dt) - omega_z = wp.to_torch(cube_object.data.root_ang_vel_w)[0, 2].item() + omega_z = cube_object.data.root_ang_vel_w.torch[0, 2].item() # No positions → force at CoM → zero torque → zero angular velocity assert abs(omega_z) < 0.01, f"Expected ~zero omega_z for force at CoM, got {omega_z}" # Should still have linear acceleration in +Y - lin_vel_y = wp.to_torch(cube_object.data.root_lin_vel_w)[0, 1].item() + lin_vel_y = cube_object.data.root_lin_vel_w.torch[0, 1].item() assert lin_vel_y > 0.1, f"Expected positive Y velocity from applied force, got {lin_vel_y}" @@ -549,7 +549,7 @@ def test_multi_cube_different_torques_from_same_force(device): body_ids, _ = cube_object.find_bodies(".*") # Position cubes: Cube 0 at (-1, 0, 1), Cube 1 at (+1, 0, 1) - root_state = wp.to_torch(cube_object.data.root_state_w).clone() + root_state = cube_object.data.root_state_w.torch.clone() root_state[0, 0] = -1.0 root_state[0, 1] = 0.0 root_state[0, 2] = 1.0 @@ -587,16 +587,16 @@ def test_multi_cube_different_torques_from_same_force(device): cube_object.update(sim.cfg.dt) # Cube 0: omega_z > 0 (force point is to the right of CoM) - omega_z_0 = wp.to_torch(cube_object.data.root_ang_vel_w)[0, 2].item() + omega_z_0 = cube_object.data.root_ang_vel_w.torch[0, 2].item() assert omega_z_0 > 0.1, f"Cube 0: expected positive omega_z, got {omega_z_0}" # Cube 1: omega_z < 0 (force point is to the left of CoM) - omega_z_1 = wp.to_torch(cube_object.data.root_ang_vel_w)[1, 2].item() + omega_z_1 = cube_object.data.root_ang_vel_w.torch[1, 2].item() assert omega_z_1 < -0.1, f"Cube 1: expected negative omega_z, got {omega_z_1}" # Both cubes should have same linear velocity in +Y (same force magnitude) - lin_vel_y_0 = wp.to_torch(cube_object.data.root_lin_vel_w)[0, 1].item() - lin_vel_y_1 = wp.to_torch(cube_object.data.root_lin_vel_w)[1, 1].item() + lin_vel_y_0 = cube_object.data.root_lin_vel_w.torch[0, 1].item() + lin_vel_y_1 = cube_object.data.root_lin_vel_w.torch[1, 1].item() assert abs(lin_vel_y_0 - lin_vel_y_1) < 0.5, ( f"Both cubes should have similar Y velocity, got {lin_vel_y_0} and {lin_vel_y_1}" ) @@ -628,7 +628,7 @@ def test_global_force_torque_far_from_origin(device): body_ids, _ = cube_object.find_bodies(".*") # Position cubes: Cube 0 near origin, Cube 1 far from origin - root_state = wp.to_torch(cube_object.data.root_state_w).clone() + root_state = cube_object.data.root_state_w.torch.clone() # Cube 0 at (0, 0, 1) root_state[0, 0] = 0.0 root_state[0, 1] = 0.0 @@ -651,7 +651,7 @@ def test_global_force_torque_far_from_origin(device): torques = torch.zeros(2, len(body_ids), 3, device=device) # Positions: each cube's CoM + (1, 0, 0) - com_pos = wp.to_torch(cube_object.data.body_com_pos_w)[:, body_ids, :3].clone() + com_pos = cube_object.data.body_com_pos_w.torch[:, body_ids, :3].clone() positions = com_pos.clone() positions[..., 0] += 1.0 # +1m X offset from CoM @@ -670,8 +670,8 @@ def test_global_force_torque_far_from_origin(device): cube_object.update(sim.cfg.dt) # Both cubes should have positive omega_z (cross((1,0,0), (0,10,0)) = (0,0,10)) - omega_z_0 = wp.to_torch(cube_object.data.root_ang_vel_w)[0, 2].item() - omega_z_1 = wp.to_torch(cube_object.data.root_ang_vel_w)[1, 2].item() + omega_z_0 = cube_object.data.root_ang_vel_w.torch[0, 2].item() + omega_z_1 = cube_object.data.root_ang_vel_w.torch[1, 2].item() assert omega_z_0 > 0.1, f"Cube 0: expected positive omega_z, got {omega_z_0}" assert omega_z_1 > 0.1, f"Cube 1: expected positive omega_z, got {omega_z_1}" @@ -689,8 +689,8 @@ def test_global_force_torque_far_from_origin(device): ) # Linear velocity in +Y should also match - lin_vel_y_0 = wp.to_torch(cube_object.data.root_lin_vel_w)[0, 1].item() - lin_vel_y_1 = wp.to_torch(cube_object.data.root_lin_vel_w)[1, 1].item() + lin_vel_y_0 = cube_object.data.root_lin_vel_w.torch[0, 1].item() + lin_vel_y_1 = cube_object.data.root_lin_vel_w.torch[1, 1].item() torch.testing.assert_close( torch.tensor(lin_vel_y_0), torch.tensor(lin_vel_y_1), @@ -723,7 +723,7 @@ def test_global_force_no_position_no_rotation_large_offset(device): body_ids, _ = cube_object.find_bodies(".*") # Place cube at large X offset - root_state = wp.to_torch(cube_object.data.default_root_state).clone() + root_state = cube_object.data.default_root_state.torch.clone() root_state[0, 0] = 2000.0 # large X position root_state[0, 1] = 0.0 root_state[0, 2] = 1.0 @@ -748,14 +748,14 @@ def test_global_force_no_position_no_rotation_large_offset(device): cube_object.update(sim.cfg.dt) # Check: angular velocity should be near zero (no rotation) - ang_vel = wp.to_torch(cube_object.data.root_ang_vel_w)[0] + ang_vel = cube_object.data.root_ang_vel_w.torch[0] assert torch.allclose(ang_vel, torch.zeros(3, device=device), atol=0.01), ( f"Expected near-zero angular velocity, got {ang_vel}. " "Global force without positions should not produce torque." ) # Check: linear velocity in Y should be positive (force is in +Y) - lin_vel = wp.to_torch(cube_object.data.root_lin_vel_w)[0] + lin_vel = cube_object.data.root_lin_vel_w.torch[0] assert lin_vel[1] > 0.1, f"Expected positive Y velocity from applied force, got {lin_vel[1]}" @@ -778,7 +778,7 @@ def test_global_force_at_com_position_no_rotation_large_offset(device): body_ids, _ = cube_object.find_bodies(".*") # Place cube at large X offset - root_state = wp.to_torch(cube_object.data.default_root_state).clone() + root_state = cube_object.data.default_root_state.torch.clone() root_state[0, 0] = 2000.0 root_state[0, 1] = 0.0 root_state[0, 2] = 1.0 @@ -807,11 +807,11 @@ def test_global_force_at_com_position_no_rotation_large_offset(device): cube_object.update(sim.cfg.dt) # Force at CoM → no rotation - ang_vel = wp.to_torch(cube_object.data.root_ang_vel_w)[0] + ang_vel = cube_object.data.root_ang_vel_w.torch[0] assert torch.allclose(ang_vel, torch.zeros(3, device=device), atol=0.01), ( f"Expected near-zero angular velocity, got {ang_vel}. " "Global force at CoM position should not produce torque." ) - lin_vel = wp.to_torch(cube_object.data.root_lin_vel_w)[0] + lin_vel = cube_object.data.root_lin_vel_w.torch[0] assert lin_vel[1] > 0.1, f"Expected positive Y velocity from applied force, got {lin_vel[1]}" diff --git a/source/isaaclab/test/utils/test_wrench_composer_vs_physx.py b/source/isaaclab/test/utils/test_wrench_composer_vs_physx.py index bb0a69132890..8d47b003369c 100644 --- a/source/isaaclab/test/utils/test_wrench_composer_vs_physx.py +++ b/source/isaaclab/test/utils/test_wrench_composer_vs_physx.py @@ -139,20 +139,20 @@ def test_composer_vs_physx_local_force(device): # Compare velocities torch.testing.assert_close( - wp.to_torch(cube_composer.data.root_lin_vel_w), - wp.to_torch(cube_raw.data.root_lin_vel_w), + cube_composer.data.root_lin_vel_w.torch, + cube_raw.data.root_lin_vel_w.torch, rtol=1e-4, atol=1e-4, ) # Both should have ~zero angular velocity (force at CoM, no torque) torch.testing.assert_close( - wp.to_torch(cube_composer.data.root_ang_vel_w), + cube_composer.data.root_ang_vel_w.torch, torch.zeros(1, 3, device=device), rtol=0.0, atol=1e-4, ) torch.testing.assert_close( - wp.to_torch(cube_raw.data.root_ang_vel_w), + cube_raw.data.root_ang_vel_w.torch, torch.zeros(1, 3, device=device), rtol=0.0, atol=1e-4, @@ -204,27 +204,27 @@ def test_composer_vs_physx_global_force(device): # Linear velocities should match (same global force, same mass) torch.testing.assert_close( - wp.to_torch(cube_composer.data.root_lin_vel_w), - wp.to_torch(cube_raw.data.root_lin_vel_w), + cube_composer.data.root_lin_vel_w.torch, + cube_raw.data.root_lin_vel_w.torch, rtol=1e-4, atol=1e-4, ) # Angular velocities should match torch.testing.assert_close( - wp.to_torch(cube_composer.data.root_ang_vel_w), - wp.to_torch(cube_raw.data.root_ang_vel_w), + cube_composer.data.root_ang_vel_w.torch, + cube_raw.data.root_ang_vel_w.torch, rtol=1e-4, atol=1e-4, ) # Both should have ~zero angular velocity (force at CoM, no torque) torch.testing.assert_close( - wp.to_torch(cube_composer.data.root_ang_vel_w), + cube_composer.data.root_ang_vel_w.torch, torch.zeros(1, 3, device=device), rtol=0.0, atol=1e-4, ) torch.testing.assert_close( - wp.to_torch(cube_raw.data.root_ang_vel_w), + cube_raw.data.root_ang_vel_w.torch, torch.zeros(1, 3, device=device), rtol=0.0, atol=1e-4, @@ -281,20 +281,20 @@ def test_composer_vs_physx_local_force_at_position(device): # Both linear and angular velocities should match torch.testing.assert_close( - wp.to_torch(cube_composer.data.root_lin_vel_w), - wp.to_torch(cube_raw.data.root_lin_vel_w), + cube_composer.data.root_lin_vel_w.torch, + cube_raw.data.root_lin_vel_w.torch, rtol=1e-4, atol=1e-4, ) torch.testing.assert_close( - wp.to_torch(cube_composer.data.root_ang_vel_w), - wp.to_torch(cube_raw.data.root_ang_vel_w), + cube_composer.data.root_ang_vel_w.torch, + cube_raw.data.root_ang_vel_w.torch, rtol=1e-4, atol=1e-4, ) # Sanity: angular velocity should be nonzero (cross-product torque) - assert torch.abs(wp.to_torch(cube_composer.data.root_ang_vel_w)[0, 2]).item() > 0.1, ( + assert torch.abs(cube_composer.data.root_ang_vel_w.torch[0, 2]).item() > 0.1, ( "Expected nonzero Z angular velocity from cross-product torque" ) @@ -319,8 +319,8 @@ def test_composer_vs_physx_global_force_at_position(device): offset = torch.zeros(1, len(body_ids), 3, device=device) offset[..., 1] = 1.0 # +1m Y offset in world frame - pos_composer = wp.to_torch(cube_composer.data.body_com_pos_w)[:, body_ids, :3].clone() + offset - pos_raw = wp.to_torch(cube_raw.data.body_com_pos_w)[:, body_ids, :3].clone() + offset + pos_composer = cube_composer.data.body_com_pos_w.torch[:, body_ids, :3].clone() + offset + pos_raw = cube_raw.data.body_com_pos_w.torch[:, body_ids, :3].clone() + offset cube_composer.permanent_wrench_composer.set_forces_and_torques( forces=forces, @@ -353,20 +353,20 @@ def test_composer_vs_physx_global_force_at_position(device): # Both linear and angular velocities should match torch.testing.assert_close( - wp.to_torch(cube_composer.data.root_lin_vel_w), - wp.to_torch(cube_raw.data.root_lin_vel_w), + cube_composer.data.root_lin_vel_w.torch, + cube_raw.data.root_lin_vel_w.torch, rtol=1e-4, atol=1e-4, ) torch.testing.assert_close( - wp.to_torch(cube_composer.data.root_ang_vel_w), - wp.to_torch(cube_raw.data.root_ang_vel_w), + cube_composer.data.root_ang_vel_w.torch, + cube_raw.data.root_ang_vel_w.torch, rtol=1e-4, atol=1e-4, ) # Sanity: angular velocity should be nonzero (cross-product torque) - assert torch.abs(wp.to_torch(cube_composer.data.root_ang_vel_w)[0, 2]).item() > 0.1, ( + assert torch.abs(cube_composer.data.root_ang_vel_w.torch[0, 2]).item() > 0.1, ( "Expected nonzero Z angular velocity from positional torque" ) @@ -416,20 +416,20 @@ def test_composer_vs_physx_local_torque(device): # Angular velocities should match torch.testing.assert_close( - wp.to_torch(cube_composer.data.root_ang_vel_w), - wp.to_torch(cube_raw.data.root_ang_vel_w), + cube_composer.data.root_ang_vel_w.torch, + cube_raw.data.root_ang_vel_w.torch, rtol=1e-4, atol=1e-4, ) # Linear velocity should be ~zero for both (no force) torch.testing.assert_close( - wp.to_torch(cube_composer.data.root_lin_vel_w), + cube_composer.data.root_lin_vel_w.torch, torch.zeros(1, 3, device=device), rtol=0.0, atol=1e-4, ) torch.testing.assert_close( - wp.to_torch(cube_raw.data.root_lin_vel_w), + cube_raw.data.root_lin_vel_w.torch, torch.zeros(1, 3, device=device), rtol=0.0, atol=1e-4, @@ -481,8 +481,8 @@ def test_composer_vs_physx_global_torque(device): # Angular velocities should match torch.testing.assert_close( - wp.to_torch(cube_composer.data.root_ang_vel_w), - wp.to_torch(cube_raw.data.root_ang_vel_w), + cube_composer.data.root_ang_vel_w.torch, + cube_raw.data.root_ang_vel_w.torch, rtol=1e-4, atol=1e-4, ) @@ -542,21 +542,21 @@ def test_composer_vs_physx_global_force_multi_env(device): # Linear velocities should match across all envs torch.testing.assert_close( - wp.to_torch(cube_composer.data.root_lin_vel_w), - wp.to_torch(cube_raw.data.root_lin_vel_w), + cube_composer.data.root_lin_vel_w.torch, + cube_raw.data.root_lin_vel_w.torch, rtol=1e-4, atol=1e-4, ) # Angular velocities should match torch.testing.assert_close( - wp.to_torch(cube_composer.data.root_ang_vel_w), - wp.to_torch(cube_raw.data.root_ang_vel_w), + cube_composer.data.root_ang_vel_w.torch, + cube_raw.data.root_ang_vel_w.torch, rtol=1e-4, atol=1e-4, ) # All envs should have ~zero angular velocity torch.testing.assert_close( - wp.to_torch(cube_composer.data.root_ang_vel_w), + cube_composer.data.root_ang_vel_w.torch, torch.zeros(NUM_CUBES_MULTI, 3, device=device), rtol=0.0, atol=1e-4, @@ -583,17 +583,17 @@ def test_composer_vs_physx_global_force_with_reset(device): cube_raw.update(sim.cfg.dt) initial_state_composer = torch.cat( [ - wp.to_torch(cube_composer.data.root_link_pos_w), - wp.to_torch(cube_composer.data.root_link_quat_w), - wp.to_torch(cube_composer.data.root_com_vel_w), + cube_composer.data.root_link_pos_w.torch, + cube_composer.data.root_link_quat_w.torch, + cube_composer.data.root_com_vel_w.torch, ], dim=-1, ).clone() initial_state_raw = torch.cat( [ - wp.to_torch(cube_raw.data.root_link_pos_w), - wp.to_torch(cube_raw.data.root_link_quat_w), - wp.to_torch(cube_raw.data.root_com_vel_w), + cube_raw.data.root_link_pos_w.torch, + cube_raw.data.root_link_quat_w.torch, + cube_raw.data.root_com_vel_w.torch, ], dim=-1, ).clone() @@ -667,20 +667,20 @@ def apply_global_force(): # All envs: composer vs raw should match torch.testing.assert_close( - wp.to_torch(cube_composer.data.root_lin_vel_w), - wp.to_torch(cube_raw.data.root_lin_vel_w), + cube_composer.data.root_lin_vel_w.torch, + cube_raw.data.root_lin_vel_w.torch, rtol=1e-4, atol=1e-4, ) torch.testing.assert_close( - wp.to_torch(cube_composer.data.root_ang_vel_w), - wp.to_torch(cube_raw.data.root_ang_vel_w), + cube_composer.data.root_ang_vel_w.torch, + cube_raw.data.root_ang_vel_w.torch, rtol=1e-4, atol=1e-4, ) # All envs should have ~zero angular velocity torch.testing.assert_close( - wp.to_torch(cube_composer.data.root_ang_vel_w), + cube_composer.data.root_ang_vel_w.torch, torch.zeros(NUM_CUBES_MULTI, 3, device=device), rtol=0.0, atol=1e-4, @@ -707,8 +707,8 @@ def test_composer_vs_physx_payload_scenario(device): cube_raw.update(sim.cfg.dt) # Record initial positions to compare displacements (cubes spawn at different Y) - init_pos_composer = wp.to_torch(cube_composer.data.root_pos_w).clone() - init_pos_raw = wp.to_torch(cube_raw.data.root_pos_w).clone() + init_pos_composer = cube_composer.data.root_pos_w.torch.clone() + init_pos_raw = cube_raw.data.root_pos_w.torch.clone() body_ids, _ = cube_composer.find_bodies(".*") @@ -744,19 +744,19 @@ def test_composer_vs_physx_payload_scenario(device): cube_raw.update(sim.cfg.dt) # Compare displacements (not absolute positions — cubes have different spawn Y) - disp_composer = wp.to_torch(cube_composer.data.root_pos_w) - init_pos_composer - disp_raw = wp.to_torch(cube_raw.data.root_pos_w) - init_pos_raw + disp_composer = cube_composer.data.root_pos_w.torch - init_pos_composer + disp_raw = cube_raw.data.root_pos_w.torch - init_pos_raw torch.testing.assert_close(disp_composer, disp_raw, rtol=1e-4, atol=1e-4) torch.testing.assert_close( - wp.to_torch(cube_composer.data.root_lin_vel_w), - wp.to_torch(cube_raw.data.root_lin_vel_w), + cube_composer.data.root_lin_vel_w.torch, + cube_raw.data.root_lin_vel_w.torch, rtol=1e-4, atol=1e-4, ) torch.testing.assert_close( - wp.to_torch(cube_composer.data.root_ang_vel_w), - wp.to_torch(cube_raw.data.root_ang_vel_w), + cube_composer.data.root_ang_vel_w.torch, + cube_raw.data.root_ang_vel_w.torch, rtol=1e-4, atol=1e-4, ) @@ -787,8 +787,8 @@ def test_composer_vs_physx_permanent_global_force_at_position_long_run(device): offset = torch.zeros(1, len(body_ids), 3, device=device) offset[..., 1] = 1.0 - pos_composer = wp.to_torch(cube_composer.data.body_com_pos_w)[:, body_ids, :3].clone() + offset - pos_raw = wp.to_torch(cube_raw.data.body_com_pos_w)[:, body_ids, :3].clone() + offset + pos_composer = cube_composer.data.body_com_pos_w.torch[:, body_ids, :3].clone() + offset + pos_raw = cube_raw.data.body_com_pos_w.torch[:, body_ids, :3].clone() + offset cube_composer.permanent_wrench_composer.set_forces_and_torques( forces=forces, @@ -819,19 +819,19 @@ def test_composer_vs_physx_permanent_global_force_at_position_long_run(device): cube_raw.update(sim.cfg.dt) torch.testing.assert_close( - wp.to_torch(cube_composer.data.root_lin_vel_w), - wp.to_torch(cube_raw.data.root_lin_vel_w), + cube_composer.data.root_lin_vel_w.torch, + cube_raw.data.root_lin_vel_w.torch, rtol=1e-3, atol=1e-3, ) torch.testing.assert_close( - wp.to_torch(cube_composer.data.root_ang_vel_w), - wp.to_torch(cube_raw.data.root_ang_vel_w), + cube_composer.data.root_ang_vel_w.torch, + cube_raw.data.root_ang_vel_w.torch, rtol=1e-3, atol=1e-3, ) # Sanity: angular velocity should be nonzero - assert torch.abs(wp.to_torch(cube_composer.data.root_ang_vel_w)).max().item() > 0.1, ( + assert torch.abs(cube_composer.data.root_ang_vel_w.torch).max().item() > 0.1, ( "Expected nonzero angular velocity from positional torque over 100 steps" ) diff --git a/source/isaaclab/test/utils/warp/test_proxy_array.py b/source/isaaclab/test/utils/warp/test_proxy_array.py new file mode 100644 index 000000000000..08bdd4bb916e --- /dev/null +++ b/source/isaaclab/test/utils/warp/test_proxy_array.py @@ -0,0 +1,550 @@ +# 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 the ProxyArray class.""" + +import warnings + +import pytest +import torch +import warp as wp + +wp.config.quiet = True +wp.init() + + +@pytest.fixture(params=["cpu", "cuda:0"]) +def device(request): + """Parametrize tests across CPU and CUDA devices.""" + return request.param + + +class TestProxyArrayBasic: + """Tests for basic ProxyArray functionality.""" + + def test_warp_returns_original(self, device): + """Test that .warp returns the original warp array.""" + from isaaclab.utils.warp.proxy_array import ProxyArray + + arr = wp.zeros(10, dtype=wp.float32, device=device) + ta = ProxyArray(arr) + assert ta.warp is arr + + def test_torch_returns_tensor(self, device): + """Test that .torch returns a torch.Tensor.""" + from isaaclab.utils.warp.proxy_array import ProxyArray + + arr = wp.zeros(10, dtype=wp.float32, device=device) + ta = ProxyArray(arr) + assert isinstance(ta.torch, torch.Tensor) + + def test_torch_is_cached(self, device): + """Test that .torch returns the same tensor object on repeated access.""" + from isaaclab.utils.warp.proxy_array import ProxyArray + + arr = wp.zeros(10, dtype=wp.float32, device=device) + ta = ProxyArray(arr) + t1 = ta.torch + t2 = ta.torch + assert t1 is t2 + + def test_torch_shares_memory(self, device): + """Test that .torch provides a zero-copy view (shares memory with warp).""" + from isaaclab.utils.warp.proxy_array import ProxyArray + + arr = wp.zeros(10, dtype=wp.float32, device=device) + ta = ProxyArray(arr) + t = ta.torch + # Modify the torch tensor + t[0] = 42.0 + # The change should be visible in the warp array + arr_np = arr.numpy() + assert arr_np[0] == 42.0 + + def test_immutable_warp_cannot_be_reassigned(self, device): + """ProxyArray._warp cannot be reassigned; callers must construct a new wrapper.""" + from isaaclab.utils.warp.proxy_array import ProxyArray + + arr = wp.zeros(10, dtype=wp.float32, device=device) + ta = ProxyArray(arr) + + with pytest.raises(AttributeError, match="immutable"): + ta._warp = wp.ones(10, dtype=wp.float32, device=device) + with pytest.raises(AttributeError, match="immutable"): + ta.new_field = 42 # arbitrary attribute writes also blocked + + def test_immutable_allows_internal_torch_cache(self, device): + """Lazy .torch caching still works — only _torch_cache is allowed as a post-init write.""" + from isaaclab.utils.warp.proxy_array import ProxyArray + + arr = wp.zeros(10, dtype=wp.float32, device=device) + ta = ProxyArray(arr) + # First access populates the cache; no exception. + first = ta.torch + # Subsequent accesses return the same cached tensor. + second = ta.torch + assert first is second + + def test_cuda_array_interface(self): + """Test that __cuda_array_interface__ delegates to the underlying warp array.""" + from isaaclab.utils.warp.proxy_array import ProxyArray + + arr = wp.zeros(10, dtype=wp.float32, device="cuda:0") + ta = ProxyArray(arr) + cai = ta.__cuda_array_interface__ + assert isinstance(cai, dict) + assert "data" in cai + assert "shape" in cai + assert cai["shape"] == arr.__cuda_array_interface__["shape"] + + def test_cuda_array_interface_not_on_cpu(self): + """Test that __cuda_array_interface__ raises AttributeError on CPU arrays.""" + from isaaclab.utils.warp.proxy_array import ProxyArray + + arr = wp.zeros(10, dtype=wp.float32, device="cpu") + ta = ProxyArray(arr) + with pytest.raises(AttributeError): + _ = ta.__cuda_array_interface__ + + def test_wp_launch_accepts_proxy_array(self): + """Test that wp.launch() can consume a ProxyArray via __cuda_array_interface__.""" + from isaaclab.utils.warp.proxy_array import ProxyArray + + @wp.kernel + def _add_one(src: wp.array(dtype=wp.float32), dst: wp.array(dtype=wp.float32)): + i = wp.tid() + dst[i] = src[i] + 1.0 + + src = ProxyArray(wp.zeros(5, dtype=wp.float32, device="cuda:0")) + dst = ProxyArray(wp.zeros(5, dtype=wp.float32, device="cuda:0")) + wp.launch(_add_one, dim=5, inputs=[src], outputs=[dst], device="cuda:0") + wp.synchronize_device("cuda:0") + assert dst.torch[0].item() == 1.0 + assert dst.torch[4].item() == 1.0 + + +class TestProxyArrayStructuredTypes: + """Tests for ProxyArray with structured warp types (vec3f, quatf, etc).""" + + def test_vec3f_shape(self, device): + """Test that vec3f arrays produce (N, 3) torch tensors.""" + from isaaclab.utils.warp.proxy_array import ProxyArray + + arr = wp.zeros(8, dtype=wp.vec3f, device=device) + ta = ProxyArray(arr) + assert ta.torch.shape == (8, 3) + + def test_quatf_shape(self, device): + """Test that quatf arrays produce (N, 4) torch tensors.""" + from isaaclab.utils.warp.proxy_array import ProxyArray + + arr = wp.zeros(8, dtype=wp.quatf, device=device) + ta = ProxyArray(arr) + assert ta.torch.shape == (8, 4) + + def test_transformf_shape(self, device): + """Test that transformf arrays produce (N, 7) torch tensors.""" + from isaaclab.utils.warp.proxy_array import ProxyArray + + arr = wp.zeros(8, dtype=wp.transformf, device=device) + ta = ProxyArray(arr) + assert ta.torch.shape == (8, 7) + + def test_spatial_vectorf_shape(self, device): + """Test that spatial_vectorf arrays produce (N, 6) torch tensors.""" + from isaaclab.utils.warp.proxy_array import ProxyArray + + arr = wp.zeros(8, dtype=wp.spatial_vectorf, device=device) + ta = ProxyArray(arr) + assert ta.torch.shape == (8, 6) + + def test_2d_vec3f_shape(self, device): + """Test that 2D vec3f arrays produce (N, M, 3) torch tensors.""" + from isaaclab.utils.warp.proxy_array import ProxyArray + + arr = wp.zeros((4, 5), dtype=wp.vec3f, device=device) + ta = ProxyArray(arr) + assert ta.torch.shape == (4, 5, 3) + + +class TestProxyArrayQuatfTorchAccessWarning: + """Tests for the WARN_ON_TORCH_QUATF_ACCESS opt-in runtime detector.""" + + def test_default_no_warning(self, device, monkeypatch): + """No env var → quatf .torch access is silent.""" + from isaaclab.utils.warp.proxy_array import ProxyArray + + monkeypatch.delenv("WARN_ON_TORCH_QUATF_ACCESS", raising=False) + ta = ProxyArray(wp.zeros(4, dtype=wp.quatf, device=device)) + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + _ = ta.torch + assert [x for x in w if issubclass(x.category, UserWarning)] == [] + + def test_env_set_warns_on_quatf(self, device, monkeypatch): + """WARN_ON_TORCH_QUATF_ACCESS=1 → quatf .torch read emits a UserWarning at the call site.""" + from isaaclab.utils.warp.proxy_array import ProxyArray + + monkeypatch.setenv("WARN_ON_TORCH_QUATF_ACCESS", "1") + ta = ProxyArray(wp.zeros(4, dtype=wp.quatf, device=device)) + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + _ = ta.torch # the .torch read on this line is what should be reported + user_warns = [x for x in w if issubclass(x.category, UserWarning)] + assert len(user_warns) == 1 + assert "quatf" in str(user_warns[0].message) + assert "(w, x, y, z)" in str(user_warns[0].message) + assert "(x, y, z, w)" in str(user_warns[0].message) + # stacklevel=2 → the warning's filename is this test file, not proxy_array.py + assert user_warns[0].filename == __file__ + + def test_env_set_does_not_warn_on_non_quatf(self, device, monkeypatch): + """The detector only fires for wp.quatf — float32 / vec3f / transformf are silent.""" + from isaaclab.utils.warp.proxy_array import ProxyArray + + monkeypatch.setenv("WARN_ON_TORCH_QUATF_ACCESS", "1") + for dtype in (wp.float32, wp.vec3f, wp.transformf, wp.spatial_vectorf): + ta = ProxyArray(wp.zeros(4, dtype=dtype, device=device)) + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + _ = ta.torch + assert [x for x in w if issubclass(x.category, UserWarning)] == [] + + def test_env_zero_does_not_warn(self, device, monkeypatch): + """WARN_ON_TORCH_QUATF_ACCESS=0 → silent (only ``"1"`` enables the detector).""" + from isaaclab.utils.warp.proxy_array import ProxyArray + + monkeypatch.setenv("WARN_ON_TORCH_QUATF_ACCESS", "0") + ta = ProxyArray(wp.zeros(4, dtype=wp.quatf, device=device)) + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + _ = ta.torch + assert [x for x in w if issubclass(x.category, UserWarning)] == [] + + +class TestProxyArrayConvenienceProperties: + """Tests for convenience properties: shape, dtype, device, len, repr.""" + + def test_shape(self, device): + """Test that .shape returns the warp array shape.""" + from isaaclab.utils.warp.proxy_array import ProxyArray + + arr = wp.zeros((3, 4), dtype=wp.float32, device=device) + ta = ProxyArray(arr) + assert ta.shape == (3, 4) + + def test_dtype(self, device): + """Test that .dtype returns the warp dtype.""" + from isaaclab.utils.warp.proxy_array import ProxyArray + + arr = wp.zeros(10, dtype=wp.float32, device=device) + ta = ProxyArray(arr) + assert ta.dtype == wp.float32 + + def test_device(self, device): + """Test that .device returns the warp device string.""" + from isaaclab.utils.warp.proxy_array import ProxyArray + + arr = wp.zeros(10, dtype=wp.float32, device=device) + ta = ProxyArray(arr) + assert ta.device == arr.device + + def test_len(self, device): + """Test that len() returns the first dimension size.""" + from isaaclab.utils.warp.proxy_array import ProxyArray + + arr = wp.zeros((7, 3), dtype=wp.float32, device=device) + ta = ProxyArray(arr) + assert len(ta) == 7 + + def test_repr(self, device): + """Test that repr() contains ProxyArray and key info.""" + from isaaclab.utils.warp.proxy_array import ProxyArray + + arr = wp.zeros(5, dtype=wp.float32, device=device) + ta = ProxyArray(arr) + r = repr(ta) + assert "ProxyArray" in r + assert "float32" in r + + +class TestProxyArrayDeprecationBridge: + """Tests for the deprecation bridge: __torch_function__, operators.""" + + def setup_method(self): + """Reset the deprecation warning flag before each test.""" + from isaaclab.utils.warp.proxy_array import ProxyArray + + ProxyArray._deprecation_warned = False + + def test_torch_function_works_and_warns(self, device): + """Test that __torch_function__ enables torch ops and emits a deprecation warning.""" + from isaaclab.utils.warp.proxy_array import ProxyArray + + arr = wp.ones(5, dtype=wp.float32, device=device) + ta = ProxyArray(arr) + + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + result = torch.sum(ta) + assert len(w) == 1 + assert issubclass(w[0].category, DeprecationWarning) + assert isinstance(result, torch.Tensor) + assert result.item() == pytest.approx(5.0) + + def test_torch_cat_works_and_warns(self, device): + """Test that torch.cat works with ProxyArray and emits a deprecation warning.""" + from isaaclab.utils.warp.proxy_array import ProxyArray + + a1 = wp.ones(3, dtype=wp.float32, device=device) + a2 = wp.ones(4, dtype=wp.float32, device=device) + ta1, ta2 = ProxyArray(a1), ProxyArray(a2) + + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + result = torch.cat([ta1, ta2]) + assert len(w) == 1 + assert issubclass(w[0].category, DeprecationWarning) + assert result.shape == (7,) + + def test_arithmetic_operators_work_and_warn(self, device): + """Test that arithmetic operators work and emit deprecation warnings.""" + from isaaclab.utils.warp.proxy_array import ProxyArray + + arr = wp.ones(5, dtype=wp.float32, device=device) + ta = ProxyArray(arr) + + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + result = ta + 1.0 + assert len(w) == 1 + assert issubclass(w[0].category, DeprecationWarning) + assert isinstance(result, torch.Tensor) + expected = torch.full((5,), 2.0, device=device) + torch.testing.assert_close(result, expected) + + def test_warns_only_once(self, device): + """Test that the deprecation warning is emitted only once per session.""" + from isaaclab.utils.warp.proxy_array import ProxyArray + + arr = wp.ones(5, dtype=wp.float32, device=device) + ta = ProxyArray(arr) + + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + _ = ta + 1.0 + _ = ta * 2.0 + _ = ta - 0.5 + # Only one warning despite three operations + assert len(w) == 1 + + def test_tensor_plus_proxy_array(self, device): + """Test that torch.Tensor + ProxyArray works via __torch_function__.""" + from isaaclab.utils.warp.proxy_array import ProxyArray + + arr = wp.ones(5, dtype=wp.float32, device=device) + ta = ProxyArray(arr) + t = torch.ones(5, device=device) * 2.0 + + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + result = t + ta + assert len(w) == 1 + assert issubclass(w[0].category, DeprecationWarning) + expected = torch.full((5,), 3.0, device=device) + torch.testing.assert_close(result, expected) + + @pytest.mark.parametrize( + "op, scalar, expected", + [ + ("+", 1.0, [2.0, 3.0]), + ("-", 1.0, [0.0, 1.0]), + ("*", 2.0, [2.0, 4.0]), + ("/", 2.0, [0.5, 1.0]), + ("**", 2.0, [1.0, 4.0]), + ], + ) + def test_binary_operators(self, op, scalar, expected): + """Test forward binary operators: +, -, *, /, **.""" + from isaaclab.utils.warp.proxy_array import ProxyArray + + ProxyArray._deprecation_warned = True + ta = ProxyArray(wp.array([1.0, 2.0], dtype=wp.float32, device="cpu")) # noqa: F841 + result = eval(f"ta {op} scalar") # noqa: S307 + assert torch.allclose(result, torch.tensor(expected)) + + @pytest.mark.parametrize( + "op, scalar, expected", + [ + ("+", 1.0, [2.0, 3.0]), + ("-", 1.0, [0.0, -1.0]), + ("*", 2.0, [2.0, 4.0]), + ("/", 2.0, [2.0, 1.0]), + ("**", 2.0, [2.0, 4.0]), + ], + ) + def test_reflected_operators(self, op, scalar, expected): + """Test reflected binary operators: scalar op ProxyArray.""" + from isaaclab.utils.warp.proxy_array import ProxyArray + + ProxyArray._deprecation_warned = True + ta = ProxyArray(wp.array([1.0, 2.0], dtype=wp.float32, device="cpu")) # noqa: F841 + result = eval(f"scalar {op} ta") # noqa: S307 + assert torch.allclose(result, torch.tensor(expected)) + + def test_proxy_array_op_proxy_array(self): + """Test binary operations between two ProxyArray instances.""" + from isaaclab.utils.warp.proxy_array import ProxyArray + + ProxyArray._deprecation_warned = True + ta1 = ProxyArray(wp.array([1.0, 2.0], dtype=wp.float32, device="cpu")) + ta2 = ProxyArray(wp.array([3.0, 4.0], dtype=wp.float32, device="cpu")) + assert torch.allclose(ta1 + ta2, torch.tensor([4.0, 6.0])) + assert torch.allclose(ta1 * ta2, torch.tensor([3.0, 8.0])) + assert torch.allclose(ta2 - ta1, torch.tensor([2.0, 2.0])) + + @pytest.mark.parametrize( + "op, values, expected", + [ + ("-", [1.0, -2.0], [-1.0, 2.0]), + ("+", [1.0, -2.0], [1.0, -2.0]), + ("abs", [-1.0, 2.0], [1.0, 2.0]), + ], + ) + def test_unary_operators(self, op, values, expected): + """Test unary operators: -, +, abs.""" + from isaaclab.utils.warp.proxy_array import ProxyArray + + ProxyArray._deprecation_warned = True + ta = ProxyArray(wp.array(values, dtype=wp.float32, device="cpu")) # noqa: F841 + result = eval(f"{op}(ta)" if op == "abs" else f"{op}ta") # noqa: S307 + assert torch.allclose(result, torch.tensor(expected)) + + @pytest.mark.parametrize( + "op, expected", + [ + ("==", [False, True, False]), + ("!=", [True, False, True]), + ("<", [True, False, False]), + ("<=", [True, True, False]), + (">", [False, False, True]), + (">=", [False, True, True]), + ], + ) + def test_comparison_operators(self, op, expected): + """Test comparison operators: ==, !=, <, <=, >, >=.""" + from isaaclab.utils.warp.proxy_array import ProxyArray + + ProxyArray._deprecation_warned = True + ta = ProxyArray(wp.array([1.0, 2.0, 3.0], dtype=wp.float32, device="cpu")) # noqa: F841 + result = eval(f"ta {op} 2.0") # noqa: S307 + assert result.tolist() == expected + + def test_getitem_1d(self): + """Test 1D indexing via __getitem__.""" + from isaaclab.utils.warp.proxy_array import ProxyArray + + ProxyArray._deprecation_warned = True + ta = ProxyArray(wp.array([10.0, 20.0, 30.0], dtype=wp.float32, device="cpu")) + assert ta[0].item() == 10.0 + assert ta[-1].item() == 30.0 + assert ta[1:].tolist() == [20.0, 30.0] + + def test_getitem_nd(self): + """Test multi-dimensional indexing via __getitem__ with structured types.""" + from isaaclab.utils.warp.proxy_array import ProxyArray + + ProxyArray._deprecation_warned = True + wp_arr = wp.zeros((3, 4), dtype=wp.vec3f, device="cpu") + ta = ProxyArray(wp_arr) + # torch view is (3, 4, 3) + result = ta[:, 0, :] + assert result.shape == (3, 3) + result = ta[0, :, 2] + assert result.shape == (4,) + + def test_setitem_writes_through(self): + """Test __setitem__ writes through to shared warp memory.""" + from isaaclab.utils.warp.proxy_array import ProxyArray + + ProxyArray._deprecation_warned = True + wp_arr = wp.array([1.0, 2.0, 3.0], dtype=wp.float32, device="cpu") + ta = ProxyArray(wp_arr) + ta[0] = 99.0 + assert wp_arr.numpy()[0] == 99.0 + + def test_getitem_warns(self): + """Test __getitem__ emits deprecation warning.""" + from isaaclab.utils.warp.proxy_array import ProxyArray + + ProxyArray._deprecation_warned = False + ta = ProxyArray(wp.array([1.0, 2.0], dtype=wp.float32, device="cpu")) + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + _ = ta[0] + deprecation_warnings = [x for x in w if issubclass(x.category, DeprecationWarning)] + assert len(deprecation_warnings) == 1 + + +class TestWpToTorchShim: + """Tests for the ``wp.to_torch`` shim installed by ``isaaclab.utils.warp``. + + The shim makes legacy call sites like ``wp.to_torch(asset.data.joint_pos)`` + keep working after the ProxyArray migration, instead of raising + ``AttributeError`` on ``requires_grad`` lookup. + """ + + def test_raw_wp_array_unchanged(self): + """``wp.to_torch(wp.array)`` must still produce a zero-copy torch view.""" + import isaaclab.utils.warp # noqa: F401 # ensure shim is installed + + arr = wp.array([1.0, 2.0, 3.0], dtype=wp.float32, device="cpu") + t = wp.to_torch(arr) + assert isinstance(t, torch.Tensor) + assert t.shape == (3,) + + def test_proxy_array_returns_torch_with_warning(self): + """``wp.to_torch(ProxyArray)`` returns the cached .torch view and warns once.""" + import isaaclab.utils.warp as iw # noqa: F401 # ensure shim is installed + from isaaclab.utils.warp.proxy_array import ProxyArray + + # Reset the module-level one-shot flag so the warning fires in this test. + iw._WP_TO_TORCH_WARNED = False + + arr = wp.array([7.0, 8.0], dtype=wp.float32, device="cpu") + proxy = ProxyArray(arr) + + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + t = wp.to_torch(proxy) + deprecation = [w for w in caught if issubclass(w.category, DeprecationWarning)] + + assert isinstance(t, torch.Tensor) + assert t is proxy.torch, "shim should return the cached .torch view" + assert len(deprecation) == 1 + assert "ProxyArray" in str(deprecation[0].message) + + def test_proxy_array_warning_is_one_shot(self): + """Repeated ``wp.to_torch(ProxyArray)`` calls must not spam warnings.""" + import isaaclab.utils.warp as iw + from isaaclab.utils.warp.proxy_array import ProxyArray + + iw._WP_TO_TORCH_WARNED = True # pretend the warning already fired + + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + wp.to_torch(ProxyArray(wp.zeros(2, dtype=wp.float32, device="cpu"))) + deprecation = [w for w in caught if issubclass(w.category, DeprecationWarning)] + + assert not deprecation, "shim must not re-warn after the first call" + + def test_requires_grad_forwarded_to_raw_wp_array(self): + """The ``requires_grad`` kwarg still reaches the original ``wp.to_torch``.""" + import isaaclab.utils.warp # noqa: F401 + + arr = wp.array([1.0, 2.0], dtype=wp.float32, device="cpu") + t = wp.to_torch(arr, requires_grad=False) + assert isinstance(t, torch.Tensor) + assert t.requires_grad is False diff --git a/source/isaaclab_experimental/config/extension.toml b/source/isaaclab_experimental/config/extension.toml index c172c3bee876..9bcfc0753383 100644 --- a/source/isaaclab_experimental/config/extension.toml +++ b/source/isaaclab_experimental/config/extension.toml @@ -1,7 +1,7 @@ [package] # Note: Semantic Versioning is used: https://semver.org/ -version = "0.0.2" +version = "0.0.3" # Description title = "Experimental playground for upcoming IsaacLab features" diff --git a/source/isaaclab_experimental/docs/CHANGELOG.rst b/source/isaaclab_experimental/docs/CHANGELOG.rst index c8bda39e0e91..5f19d1fb8c5a 100644 --- a/source/isaaclab_experimental/docs/CHANGELOG.rst +++ b/source/isaaclab_experimental/docs/CHANGELOG.rst @@ -1,6 +1,30 @@ Changelog --------- +0.0.3 (2026-04-27) +~~~~~~~~~~~~~~~~~~ + +Changed +^^^^^^^ + +* Updated the Warp-graphable MDP terms and the Warp inhand-manipulation env to read + asset/sensor data via the explicit :attr:`~isaaclab.utils.warp.ProxyArray.warp` + accessor when the value flows into a ``wp.launch`` call (or a sim-write helper that + forwards to one). Affected modules: + :mod:`isaaclab_experimental.envs.mdp.observations`, + :mod:`isaaclab_experimental.envs.mdp.rewards`, + :mod:`isaaclab_experimental.envs.mdp.terminations`, + :mod:`isaaclab_experimental.envs.mdp.events`, + :mod:`isaaclab_experimental.envs.mdp.actions.joint_actions`, and + :mod:`isaaclab_tasks_experimental.direct.inhand_manipulation.inhand_manipulation_warp_env`. + The previous code relied on ``ProxyArray``'s ``__cuda_array_interface__`` bridge, + which works but is not explicit. No behavior change. +* Replaced ``wp.to_torch(asset.data.joint_pos).shape[1]`` in + :class:`~isaaclab_experimental.managers.ObservationManager` with + ``asset.data.joint_pos.shape[1]`` — :class:`~isaaclab.utils.warp.ProxyArray` forwards + ``shape``, so the round-trip through ``wp.to_torch`` is no longer needed. + + 0.0.2 (2026-03-16) ~~~~~~~~~~~~~~~~~~ diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/actions/joint_actions.py b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/actions/joint_actions.py index 441bde86b4de..f379f66c1900 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/actions/joint_actions.py +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/actions/joint_actions.py @@ -272,7 +272,7 @@ def __init__(self, cfg: actions_cfg.JointPositionActionCfg, env: ManagerBasedEnv super().__init__(cfg, env) # use default joint positions as offset if cfg.use_default_offset: - defaults_np = self._asset.data.default_joint_pos.numpy() + defaults_np = self._asset.data.default_joint_pos.warp.numpy() if isinstance(self._joint_ids, slice): offset_vals = defaults_np[0, :].tolist() else: diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/events.py b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/events.py index 1e9dad36ddaa..e1732df0c026 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/events.py +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/events.py @@ -90,7 +90,7 @@ def randomize_rigid_body_com( inputs=[ env_mask, env.rng_state_wp, - asset.data.body_com_pos_b, + asset.data.body_com_pos_b.warp, asset_cfg.body_ids_wp, fn._com_lo, fn._com_hi, @@ -99,7 +99,7 @@ def randomize_rigid_body_com( ) # Notify the solver that inertial properties changed (COM position affects inertia). - asset.set_coms_mask(coms=asset.data.body_com_pos_b, env_mask=env_mask) + asset.set_coms_mask(coms=asset.data.body_com_pos_b.warp, env_mask=env_mask) # --------------------------------------------------------------------------- @@ -250,7 +250,7 @@ def push_by_setting_velocity( inputs=[ env_mask, env.rng_state_wp, - asset.data.root_vel_w, + asset.data.root_vel_w.warp, push_by_setting_velocity._scratch_vel, push_by_setting_velocity._lin_lo, push_by_setting_velocity._lin_hi, @@ -369,8 +369,8 @@ def reset_root_state_uniform( inputs=[ env_mask, env.rng_state_wp, - asset.data.default_root_pose, - asset.data.default_root_vel, + asset.data.default_root_pose.warp, + asset.data.default_root_vel.warp, env.env_origins_wp, reset_root_state_uniform._scratch_pose, reset_root_state_uniform._scratch_vel, @@ -474,12 +474,12 @@ def reset_joints_by_offset( env_mask, asset_cfg.joint_ids_wp, env.rng_state_wp, - asset.data.default_joint_pos, - asset.data.default_joint_vel, - asset.data.joint_pos, - asset.data.joint_vel, - asset.data.soft_joint_pos_limits, - asset.data.soft_joint_vel_limits, + asset.data.default_joint_pos.warp, + asset.data.default_joint_vel.warp, + asset.data.joint_pos.warp, + asset.data.joint_vel.warp, + asset.data.soft_joint_pos_limits.warp, + asset.data.soft_joint_vel_limits.warp, float(position_range[0]), float(position_range[1]), float(velocity_range[0]), @@ -489,8 +489,8 @@ def reset_joints_by_offset( ) # Sync derived buffers (_previous_joint_vel, joint_acc) for reset envs. - asset.write_joint_position_to_sim_mask(position=asset.data.joint_pos, env_mask=env_mask) - asset.write_joint_velocity_to_sim_mask(velocity=asset.data.joint_vel, env_mask=env_mask) + asset.write_joint_position_to_sim_mask(position=asset.data.joint_pos.warp, env_mask=env_mask) + asset.write_joint_velocity_to_sim_mask(velocity=asset.data.joint_vel.warp, env_mask=env_mask) @wp.kernel @@ -565,12 +565,12 @@ def reset_joints_by_scale( env_mask, asset_cfg.joint_ids_wp, env.rng_state_wp, - asset.data.default_joint_pos, - asset.data.default_joint_vel, - asset.data.joint_pos, - asset.data.joint_vel, - asset.data.soft_joint_pos_limits, - asset.data.soft_joint_vel_limits, + asset.data.default_joint_pos.warp, + asset.data.default_joint_vel.warp, + asset.data.joint_pos.warp, + asset.data.joint_vel.warp, + asset.data.soft_joint_pos_limits.warp, + asset.data.soft_joint_vel_limits.warp, float(position_range[0]), float(position_range[1]), float(velocity_range[0]), @@ -580,5 +580,5 @@ def reset_joints_by_scale( ) # Sync derived buffers (_previous_joint_vel, joint_acc) for reset envs. - asset.write_joint_position_to_sim_mask(position=asset.data.joint_pos, env_mask=env_mask) - asset.write_joint_velocity_to_sim_mask(velocity=asset.data.joint_vel, env_mask=env_mask) + asset.write_joint_position_to_sim_mask(position=asset.data.joint_pos.warp, env_mask=env_mask) + asset.write_joint_velocity_to_sim_mask(velocity=asset.data.joint_vel.warp, env_mask=env_mask) diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/observations.py b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/observations.py index 29fdedfd98b0..5760d1f03e8c 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/observations.py +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/observations.py @@ -94,7 +94,7 @@ def base_pos_z(env: ManagerBasedEnv, out, asset_cfg: SceneEntityCfg = SceneEntit wp.launch( kernel=_base_pos_z_kernel, dim=env.num_envs, - inputs=[asset.data.root_pos_w, out], + inputs=[asset.data.root_pos_w.warp, out], device=env.device, ) @@ -129,7 +129,7 @@ def base_lin_vel(env: ManagerBasedEnv, out, asset_cfg: SceneEntityCfg = SceneEnt wp.launch( kernel=_base_lin_vel_kernel, dim=env.num_envs, - inputs=[asset.data.root_link_pose_w, asset.data.root_com_vel_w, out], + inputs=[asset.data.root_link_pose_w.warp, asset.data.root_com_vel_w.warp, out], device=env.device, ) @@ -156,7 +156,7 @@ def base_ang_vel(env: ManagerBasedEnv, out, asset_cfg: SceneEntityCfg = SceneEnt wp.launch( kernel=_base_ang_vel_kernel, dim=env.num_envs, - inputs=[asset.data.root_link_pose_w, asset.data.root_com_vel_w, out], + inputs=[asset.data.root_link_pose_w.warp, asset.data.root_com_vel_w.warp, out], device=env.device, ) @@ -183,7 +183,7 @@ def projected_gravity(env: ManagerBasedEnv, out, asset_cfg: SceneEntityCfg = Sce wp.launch( kernel=_projected_gravity_kernel, dim=env.num_envs, - inputs=[asset.data.root_link_pose_w, asset.data.GRAVITY_VEC_W, out], + inputs=[asset.data.root_link_pose_w.warp, asset.data.GRAVITY_VEC_W.warp, out], device=env.device, ) @@ -208,7 +208,7 @@ def joint_pos(env: ManagerBasedEnv, out, asset_cfg: SceneEntityCfg = SceneEntity wp.launch( kernel=_joint_gather_kernel, dim=(env.num_envs, out.shape[1]), - inputs=[asset.data.joint_pos, joint_ids_wp, out], + inputs=[asset.data.joint_pos.warp, joint_ids_wp, out], device=env.device, ) @@ -244,7 +244,7 @@ def joint_pos_rel(env: ManagerBasedEnv, out, asset_cfg: SceneEntityCfg = SceneEn wp.launch( kernel=_joint_rel_gather_kernel, dim=(env.num_envs, out.shape[1]), - inputs=[asset.data.joint_pos, asset.data.default_joint_pos, joint_ids_wp, out], + inputs=[asset.data.joint_pos.warp, asset.data.default_joint_pos.warp, joint_ids_wp, out], device=env.device, ) @@ -278,7 +278,7 @@ def joint_pos_limit_normalized(env: ManagerBasedEnv, out, asset_cfg: SceneEntity wp.launch( kernel=_joint_pos_limit_normalized_kernel, dim=(env.num_envs, out.shape[1]), - inputs=[asset.data.joint_pos, asset.data.soft_joint_pos_limits, joint_ids_wp, out], + inputs=[asset.data.joint_pos.warp, asset.data.soft_joint_pos_limits.warp, joint_ids_wp, out], device=env.device, ) @@ -298,7 +298,7 @@ def joint_vel(env: ManagerBasedEnv, out, asset_cfg: SceneEntityCfg = SceneEntity wp.launch( kernel=_joint_gather_kernel, dim=(env.num_envs, out.shape[1]), - inputs=[asset.data.joint_vel, joint_ids_wp, out], + inputs=[asset.data.joint_vel.warp, joint_ids_wp, out], device=env.device, ) @@ -322,7 +322,7 @@ def joint_vel_rel(env: ManagerBasedEnv, out, asset_cfg: SceneEntityCfg = SceneEn wp.launch( kernel=_joint_rel_gather_kernel, dim=(env.num_envs, out.shape[1]), - inputs=[asset.data.joint_vel, asset.data.default_joint_vel, joint_ids_wp, out], + inputs=[asset.data.joint_vel.warp, asset.data.default_joint_vel.warp, joint_ids_wp, out], device=env.device, ) diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/rewards.py b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/rewards.py index b1a1a5ddfd93..2b98f1f9cc09 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/rewards.py +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/rewards.py @@ -99,7 +99,7 @@ def lin_vel_z_l2(env: ManagerBasedRLEnv, out, asset_cfg: SceneEntityCfg = SceneE wp.launch( kernel=_lin_vel_z_l2_kernel, dim=env.num_envs, - inputs=[asset.data.root_link_pose_w, asset.data.root_com_vel_w, out], + inputs=[asset.data.root_link_pose_w.warp, asset.data.root_com_vel_w.warp, out], device=env.device, ) @@ -121,7 +121,7 @@ def ang_vel_xy_l2(env: ManagerBasedRLEnv, out, asset_cfg: SceneEntityCfg = Scene wp.launch( kernel=_ang_vel_xy_l2_kernel, dim=env.num_envs, - inputs=[asset.data.root_link_pose_w, asset.data.root_com_vel_w, out], + inputs=[asset.data.root_link_pose_w.warp, asset.data.root_com_vel_w.warp, out], device=env.device, ) @@ -143,7 +143,7 @@ def flat_orientation_l2(env: ManagerBasedRLEnv, out, asset_cfg: SceneEntityCfg = wp.launch( kernel=_flat_orientation_l2_kernel, dim=env.num_envs, - inputs=[asset.data.root_link_pose_w, asset.data.GRAVITY_VEC_W, out], + inputs=[asset.data.root_link_pose_w.warp, asset.data.GRAVITY_VEC_W.warp, out], device=env.device, ) @@ -173,7 +173,7 @@ def joint_torques_l2(env: ManagerBasedRLEnv, out, asset_cfg: SceneEntityCfg = Sc wp.launch( kernel=_sum_sq_masked_kernel, dim=env.num_envs, - inputs=[asset.data.applied_torque, asset_cfg.joint_mask, out], + inputs=[asset.data.applied_torque.warp, asset_cfg.joint_mask, out], device=env.device, ) @@ -197,7 +197,7 @@ def joint_vel_l1(env: ManagerBasedRLEnv, out, asset_cfg: SceneEntityCfg) -> None wp.launch( kernel=_sum_abs_masked_kernel, dim=env.num_envs, - inputs=[asset.data.joint_vel, asset_cfg.joint_mask, out], + inputs=[asset.data.joint_vel.warp, asset_cfg.joint_mask, out], device=env.device, ) @@ -208,7 +208,7 @@ def joint_vel_l2(env: ManagerBasedRLEnv, out, asset_cfg: SceneEntityCfg = SceneE wp.launch( kernel=_sum_sq_masked_kernel, dim=env.num_envs, - inputs=[asset.data.joint_vel, asset_cfg.joint_mask, out], + inputs=[asset.data.joint_vel.warp, asset_cfg.joint_mask, out], device=env.device, ) @@ -219,7 +219,7 @@ def joint_acc_l2(env: ManagerBasedRLEnv, out, asset_cfg: SceneEntityCfg = SceneE wp.launch( kernel=_sum_sq_masked_kernel, dim=env.num_envs, - inputs=[asset.data.joint_acc, asset_cfg.joint_mask, out], + inputs=[asset.data.joint_acc.warp, asset_cfg.joint_mask, out], device=env.device, ) @@ -246,7 +246,7 @@ def joint_deviation_l1(env: ManagerBasedRLEnv, out, asset_cfg: SceneEntityCfg = wp.launch( kernel=_sum_abs_diff_masked_kernel, dim=env.num_envs, - inputs=[asset.data.joint_pos, asset.data.default_joint_pos, asset_cfg.joint_mask, out], + inputs=[asset.data.joint_pos.warp, asset.data.default_joint_pos.warp, asset_cfg.joint_mask, out], device=env.device, ) @@ -284,7 +284,7 @@ def joint_pos_limits(env: ManagerBasedRLEnv, out, asset_cfg: SceneEntityCfg = Sc wp.launch( kernel=_joint_pos_limits_kernel, dim=env.num_envs, - inputs=[asset.data.joint_pos, asset.data.soft_joint_pos_limits, asset_cfg.joint_mask, out], + inputs=[asset.data.joint_pos.warp, asset.data.soft_joint_pos_limits.warp, asset_cfg.joint_mask, out], device=env.device, ) @@ -376,7 +376,7 @@ def undesired_contacts(env: ManagerBasedRLEnv, out, threshold: float, sensor_cfg wp.launch( kernel=_undesired_contacts_kernel, dim=env.num_envs, - inputs=[contact_sensor.data.net_forces_w_history, sensor_cfg.body_ids_wp, threshold, out], + inputs=[contact_sensor.data.net_forces_w_history.warp, sensor_cfg.body_ids_wp, threshold, out], device=env.device, ) @@ -429,8 +429,8 @@ def track_lin_vel_xy_exp( kernel=_track_lin_vel_xy_exp_kernel, dim=env.num_envs, inputs=[ - asset.data.root_link_pose_w, - asset.data.root_com_vel_w, + asset.data.root_link_pose_w.warp, + asset.data.root_com_vel_w.warp, track_lin_vel_xy_exp._cmd_wp, 1.0 / (std * std), out, @@ -479,8 +479,8 @@ def track_ang_vel_z_exp( kernel=_track_ang_vel_z_exp_kernel, dim=env.num_envs, inputs=[ - asset.data.root_link_pose_w, - asset.data.root_com_vel_w, + asset.data.root_link_pose_w.warp, + asset.data.root_com_vel_w.warp, track_ang_vel_z_exp._cmd_wp, 2, 1.0 / (std * std), diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/terminations.py b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/terminations.py index 9ea143e4e8f0..7721469e96a4 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/terminations.py +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/terminations.py @@ -73,7 +73,7 @@ def root_height_below_minimum( wp.launch( kernel=_root_height_below_min_kernel, dim=env.num_envs, - inputs=[asset.data.root_pos_w, minimum_height, out], + inputs=[asset.data.root_pos_w.warp, minimum_height, out], device=env.device, ) @@ -109,15 +109,15 @@ def joint_pos_out_of_manual_limit( f"joint_pos_out_of_manual_limit requires SceneEntityCfg with resolved joint_mask, " f"but got None for asset '{asset_cfg.name}'." ) - if asset.data.joint_pos.shape[1] != asset_cfg.joint_mask.shape[0]: + if asset.data.joint_pos.warp.shape[1] != asset_cfg.joint_mask.shape[0]: raise ValueError( f"joint_mask length ({asset_cfg.joint_mask.shape[0]}) does not match " - f"joint_pos dim ({asset.data.joint_pos.shape[1]}) for asset '{asset_cfg.name}'." + f"joint_pos dim ({asset.data.joint_pos.warp.shape[1]}) for asset '{asset_cfg.name}'." ) wp.launch( kernel=_joint_pos_out_of_manual_limit_kernel, - dim=(env.num_envs, asset.data.joint_pos.shape[1]), - inputs=[asset.data.joint_pos, asset_cfg.joint_mask, bounds[0], bounds[1], out], + dim=(env.num_envs, asset.data.joint_pos.warp.shape[1]), + inputs=[asset.data.joint_pos.warp, asset_cfg.joint_mask, bounds[0], bounds[1], out], device=env.device, ) @@ -156,6 +156,6 @@ def illegal_contact(env: ManagerBasedRLEnv, out, threshold: float, sensor_cfg: S wp.launch( kernel=_illegal_contact_kernel, dim=env.num_envs, - inputs=[contact_sensor.data.net_forces_w_history, sensor_cfg.body_ids_wp, threshold, out], + inputs=[contact_sensor.data.net_forces_w_history.warp, sensor_cfg.body_ids_wp, threshold, out], device=env.device, ) diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/utils/io_descriptors.py b/source/isaaclab_experimental/isaaclab_experimental/envs/utils/io_descriptors.py index 19b6530b0b0e..2b72dae02aa7 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/utils/io_descriptors.py +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/utils/io_descriptors.py @@ -303,7 +303,7 @@ def record_joint_pos_offsets(output: wp.array | None, descriptor: GenericObserva ids = kwargs["asset_cfg"].joint_ids # Get the offsets of the joints for the first robot in the scene. # This assumes that all robots have the same joint offsets. - descriptor.joint_pos_offsets = wp.to_torch(asset.data.default_joint_pos).clone()[:, ids][0] + descriptor.joint_pos_offsets = asset.data.default_joint_pos.torch.clone()[:, ids][0] def record_joint_vel_offsets(output: wp.array | None, descriptor: GenericObservationIODescriptor, **kwargs): @@ -320,4 +320,4 @@ def record_joint_vel_offsets(output: wp.array | None, descriptor: GenericObserva ids = kwargs["asset_cfg"].joint_ids # Get the offsets of the joints for the first robot in the scene. # This assumes that all robots have the same joint offsets. - descriptor.joint_vel_offsets = wp.to_torch(asset.data.default_joint_vel).clone()[:, ids][0] + descriptor.joint_vel_offsets = asset.data.default_joint_vel.torch.clone()[:, ids][0] diff --git a/source/isaaclab_experimental/isaaclab_experimental/managers/observation_manager.py b/source/isaaclab_experimental/isaaclab_experimental/managers/observation_manager.py index 07287eb03fbc..74d0a9a955be 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/managers/observation_manager.py +++ b/source/isaaclab_experimental/isaaclab_experimental/managers/observation_manager.py @@ -893,7 +893,7 @@ def _infer_term_dim_scalar(self, term_cfg: ObservationTermCfg) -> int: return int(joint_ids_wp.shape[0]) joint_ids = getattr(asset_cfg, "joint_ids", slice(None)) if isinstance(joint_ids, slice): - return int(getattr(asset, "num_joints", wp.to_torch(asset.data.joint_pos).shape[1])) + return int(getattr(asset, "num_joints", asset.data.joint_pos.torch.shape[1])) return int(len(joint_ids)) def _resolve_out_dim(self, out_dim: int | str, term_cfg: ObservationTermCfg) -> int: @@ -917,7 +917,7 @@ def _resolve_out_dim(self, out_dim: int | str, term_cfg: ObservationTermCfg) -> return int(joint_ids_wp.shape[0]) joint_ids = getattr(asset_cfg, "joint_ids", slice(None)) if isinstance(joint_ids, slice): - return int(getattr(asset, "num_joints", wp.to_torch(asset.data.joint_pos).shape[1])) + return int(getattr(asset, "num_joints", asset.data.joint_pos.warp.shape[1])) return int(len(joint_ids)) if isinstance(out_dim, str) and out_dim.startswith("body:"): diff --git a/source/isaaclab_experimental/test/envs/mdp/parity_helpers.py b/source/isaaclab_experimental/test/envs/mdp/parity_helpers.py index 33a21907ee94..9aa219ee4a41 100644 --- a/source/isaaclab_experimental/test/envs/mdp/parity_helpers.py +++ b/source/isaaclab_experimental/test/envs/mdp/parity_helpers.py @@ -16,6 +16,8 @@ import torch import warp as wp +from isaaclab.utils.warp import ProxyArray + # --------------------------------------------------------------------------- # Constants (shared across all MDP parity test files) # --------------------------------------------------------------------------- @@ -64,8 +66,19 @@ def quat_rotate_inv_np(q_xyzw: np.ndarray, v: np.ndarray) -> np.ndarray: # --------------------------------------------------------------------------- -def copy_np_to_wp(dest: wp.array, src_np: np.ndarray): +def proxy_array(*args, **kwargs) -> ProxyArray: + """Build a :class:`ProxyArray` directly from ``wp.array`` constructor arguments. + + Lets the parity-test mocks read as ``self.joint_pos = proxy_array(np_data, device=...)`` + instead of ``ProxyArray(wp.array(np_data, device=...))``. + """ + return ProxyArray(wp.array(*args, **kwargs)) + + +def copy_np_to_wp(dest: wp.array | ProxyArray, src_np: np.ndarray): """In-place overwrite of a warp array's contents from numpy (preserves pointer).""" + if isinstance(dest, ProxyArray): + dest = dest.warp tmp = wp.array(src_np, dtype=dest.dtype, device=str(dest.device)) wp.copy(dest, tmp) @@ -162,55 +175,55 @@ def __init__(self, num_envs=NUM_ENVS, num_joints=NUM_JOINTS, device=DEVICE, seed rng = np.random.RandomState(seed) # --- Joint state (float32 2D) --- - self.joint_pos = wp.array(rng.randn(num_envs, num_joints).astype(np.float32), device=device) - self.joint_vel = wp.array(rng.randn(num_envs, num_joints).astype(np.float32) * 2.0, device=device) - self.joint_acc = wp.array(rng.randn(num_envs, num_joints).astype(np.float32) * 0.5, device=device) - self.default_joint_pos = wp.array(rng.randn(num_envs, num_joints).astype(np.float32) * 0.01, device=device) - self.default_joint_vel = wp.array(np.zeros((num_envs, num_joints), dtype=np.float32), device=device) - self.applied_torque = wp.array(rng.randn(num_envs, num_joints).astype(np.float32) * 10.0, device=device) - self.computed_torque = wp.array(rng.randn(num_envs, num_joints).astype(np.float32) * 10.0, device=device) + self.joint_pos = proxy_array(rng.randn(num_envs, num_joints).astype(np.float32), device=device) + self.joint_vel = proxy_array(rng.randn(num_envs, num_joints).astype(np.float32) * 2.0, device=device) + self.joint_acc = proxy_array(rng.randn(num_envs, num_joints).astype(np.float32) * 0.5, device=device) + self.default_joint_pos = proxy_array(rng.randn(num_envs, num_joints).astype(np.float32) * 0.01, device=device) + self.default_joint_vel = proxy_array(np.zeros((num_envs, num_joints), dtype=np.float32), device=device) + self.applied_torque = proxy_array(rng.randn(num_envs, num_joints).astype(np.float32) * 10.0, device=device) + self.computed_torque = proxy_array(rng.randn(num_envs, num_joints).astype(np.float32) * 10.0, device=device) # --- Soft joint limits --- limits_np = np.zeros((num_envs, num_joints, 2), dtype=np.float32) limits_np[:, :, 0] = -3.14 limits_np[:, :, 1] = 3.14 - self.soft_joint_pos_limits = wp.array(limits_np, dtype=wp.vec2f, device=device) - self.soft_joint_vel_limits = wp.array(np.full((num_envs, num_joints), 10.0, dtype=np.float32), device=device) + self.soft_joint_pos_limits = proxy_array(limits_np, dtype=wp.vec2f, device=device) + self.soft_joint_vel_limits = proxy_array(np.full((num_envs, num_joints), 10.0, dtype=np.float32), device=device) # --- Root state --- root_pos_np = rng.randn(num_envs, 3).astype(np.float32) root_pos_np[:, 2] = np.abs(root_pos_np[:, 2]) + 0.1 # positive heights - self.root_pos_w = wp.array(root_pos_np, dtype=wp.vec3f, device=device) + self.root_pos_w = proxy_array(root_pos_np, dtype=wp.vec3f, device=device) # Unit quaternions quat_np = rng.randn(num_envs, 4).astype(np.float32) quat_np /= np.linalg.norm(quat_np, axis=1, keepdims=True) - self.root_quat_w = wp.array(quat_np, dtype=wp.quatf, device=device) + self.root_quat_w = proxy_array(quat_np, dtype=wp.quatf, device=device) # Tier 1 compound: root_link_pose_w (transformf = pos + quat) pose_np = np.zeros((num_envs, 7), dtype=np.float32) pose_np[:, :3] = root_pos_np pose_np[:, 3:] = quat_np - self.root_link_pose_w = wp.array(pose_np, dtype=wp.transformf, device=device) + self.root_link_pose_w = proxy_array(pose_np, dtype=wp.transformf, device=device) # World-frame velocities lin_vel_w_np = rng.randn(num_envs, 3).astype(np.float32) ang_vel_w_np = rng.randn(num_envs, 3).astype(np.float32) - self.root_lin_vel_w = wp.array(lin_vel_w_np, dtype=wp.vec3f, device=device) - self.root_ang_vel_w = wp.array(ang_vel_w_np, dtype=wp.vec3f, device=device) + self.root_lin_vel_w = proxy_array(lin_vel_w_np, dtype=wp.vec3f, device=device) + self.root_ang_vel_w = proxy_array(ang_vel_w_np, dtype=wp.vec3f, device=device) # Tier 1 compound: root_com_vel_w (spatial_vectorf: top=linear, bottom=angular) vel_np = np.zeros((num_envs, 6), dtype=np.float32) vel_np[:, :3] = lin_vel_w_np vel_np[:, 3:] = ang_vel_w_np - self.root_com_vel_w = wp.array(vel_np, dtype=wp.spatial_vectorf, device=device) + self.root_com_vel_w = proxy_array(vel_np, dtype=wp.spatial_vectorf, device=device) # Gravity direction constant (1D array to match kernel signatures) - self.GRAVITY_VEC_W = wp.array([0.0, 0.0, -1.0], dtype=wp.vec3f, device=device) + self.GRAVITY_VEC_W = proxy_array([0.0, 0.0, -1.0], dtype=wp.vec3f, device=device) # Derived body-frame quantities (consistent with Tier 1 compounds) - self.root_lin_vel_b = wp.array(quat_rotate_inv_np(quat_np, lin_vel_w_np), dtype=wp.vec3f, device=device) - self.root_ang_vel_b = wp.array(quat_rotate_inv_np(quat_np, ang_vel_w_np), dtype=wp.vec3f, device=device) + self.root_lin_vel_b = proxy_array(quat_rotate_inv_np(quat_np, lin_vel_w_np), dtype=wp.vec3f, device=device) + self.root_ang_vel_b = proxy_array(quat_rotate_inv_np(quat_np, ang_vel_w_np), dtype=wp.vec3f, device=device) # --- projected_gravity_b and body-level data --- if num_bodies > 0: @@ -218,40 +231,42 @@ def __init__(self, num_envs=NUM_ENVS, num_joints=NUM_JOINTS, device=DEVICE, seed grav_np = rng.randn(num_envs, num_bodies, 3).astype(np.float32) grav_np[:, :, 2] = -1.0 grav_np /= np.linalg.norm(grav_np, axis=2, keepdims=True) - self.projected_gravity_b = wp.array(grav_np, dtype=wp.vec3f, device=device) + self.projected_gravity_b = proxy_array(grav_np, dtype=wp.vec3f, device=device) # body_pose_w: (num_envs, num_bodies) transformf bpose_np = np.zeros((num_envs, num_bodies, 7), dtype=np.float32) bpose_np[:, :, :3] = rng.randn(num_envs, num_bodies, 3).astype(np.float32) bpose_np[:, :, 3:7] = [0.0, 0.0, 0.0, 1.0] - self.body_pose_w = wp.array(bpose_np, dtype=wp.transformf, device=device) + self.body_pose_w = proxy_array(bpose_np, dtype=wp.transformf, device=device) # body_lin_acc_w: (num_envs, num_bodies) vec3f - self.body_lin_acc_w = wp.array( + self.body_lin_acc_w = proxy_array( rng.randn(num_envs, num_bodies, 3).astype(np.float32), dtype=wp.vec3f, device=device ) # body_com_pos_b: (num_envs, num_bodies) vec3f - self.body_com_pos_b = wp.array( + self.body_com_pos_b = proxy_array( rng.randn(num_envs, num_bodies, 3).astype(np.float32) * 0.01, dtype=wp.vec3f, device=device ) else: # Root-level projected_gravity_b: (num_envs,) vec3f — derived from root quat - self.projected_gravity_b = wp.array( + self.projected_gravity_b = proxy_array( quat_rotate_inv_np(quat_np, np.tile(GRAVITY_DIR_NP, (num_envs, 1))), dtype=wp.vec3f, device=device, ) # --- Event-specific data --- - self.root_vel_w = wp.array(rng.randn(num_envs, 6).astype(np.float32), dtype=wp.spatial_vectorf, device=device) + self.root_vel_w = proxy_array( + rng.randn(num_envs, 6).astype(np.float32), dtype=wp.spatial_vectorf, device=device + ) default_pose_np = np.zeros((num_envs, 7), dtype=np.float32) default_pose_np[:, 0:3] = rng.randn(num_envs, 3).astype(np.float32) * 0.1 default_pose_np[:, 3:7] = [0.0, 0.0, 0.0, 1.0] - self.default_root_pose = wp.array(default_pose_np, dtype=wp.transformf, device=device) + self.default_root_pose = proxy_array(default_pose_np, dtype=wp.transformf, device=device) - self.default_root_vel = wp.array( + self.default_root_vel = proxy_array( np.zeros((num_envs, 6), dtype=np.float32), dtype=wp.spatial_vectorf, device=device ) @@ -417,7 +432,7 @@ def mutate_root_state(rng: np.random.RandomState, art_data: MockArticulationData # Root-level projected_gravity_b (1D) is derived from quat. # Multi-body (2D) is mutated separately by callers. - if art_data.projected_gravity_b.ndim == 1: + if art_data.projected_gravity_b.warp.ndim == 1: copy_np_to_wp( art_data.projected_gravity_b, quat_rotate_inv_np(quat_np, np.tile(GRAVITY_DIR_NP, (num_envs, 1))), @@ -487,7 +502,7 @@ class MockContactSensorData: def __init__(self, num_envs=NUM_ENVS, num_history=NUM_HISTORY, num_bodies=NUM_BODIES, device=DEVICE, seed=77): rng = np.random.RandomState(seed) - self.net_forces_w_history = wp.array( + self.net_forces_w_history = proxy_array( rng.randn(num_envs, num_history, num_bodies, 3).astype(np.float32), dtype=wp.vec3f, device=device, diff --git a/source/isaaclab_experimental/test/envs/mdp/test_actions_warp_parity.py b/source/isaaclab_experimental/test/envs/mdp/test_actions_warp_parity.py index 237632efa84c..f2db4f64fbe9 100644 --- a/source/isaaclab_experimental/test/envs/mdp/test_actions_warp_parity.py +++ b/source/isaaclab_experimental/test/envs/mdp/test_actions_warp_parity.py @@ -121,7 +121,7 @@ def test_joint_position_default_offset(self, env, asset, art_data, actions_wp): term.apply_actions() # Processed = raw * 1.0 + default_joint_pos[0] - defaults = wp.to_torch(art_data.default_joint_pos)[0] + defaults = art_data.default_joint_pos.torch[0] raw = wp.to_torch(actions_wp) expected = raw + defaults.unsqueeze(0) assert_close(term.processed_actions, expected) @@ -211,7 +211,7 @@ def test_position_with_default_offset(self, env, art_data, actions_wp): term.process_actions(actions_wp, action_offset=0) raw = wp.to_torch(actions_wp) - defaults = wp.to_torch(art_data.default_joint_pos)[0] + defaults = art_data.default_joint_pos.torch[0] expected = raw * 1.0 + defaults.unsqueeze(0) assert_close(term.processed_actions, expected) @@ -222,6 +222,6 @@ def test_position_scale_and_offset(self, env, art_data, actions_wp): term.process_actions(actions_wp, action_offset=0) raw = wp.to_torch(actions_wp) - defaults = wp.to_torch(art_data.default_joint_pos)[0] + defaults = art_data.default_joint_pos.torch[0] expected = raw * 2.0 + defaults.unsqueeze(0) assert_close(term.processed_actions, expected) diff --git a/source/isaaclab_experimental/test/envs/mdp/test_events_warp_parity.py b/source/isaaclab_experimental/test/envs/mdp/test_events_warp_parity.py index d29f15f52d8b..24b9eb022665 100644 --- a/source/isaaclab_experimental/test/envs/mdp/test_events_warp_parity.py +++ b/source/isaaclab_experimental/test/envs/mdp/test_events_warp_parity.py @@ -157,16 +157,16 @@ def test_reset_joints_by_offset_parity(self, warp_env, stable_env, art_data, all warp_env, mask, position_range=(0.0, 0.0), velocity_range=(0.0, 0.0), asset_cfg=cfg ) wp.synchronize() - warp_pos = wp.to_torch(art_data.joint_pos).clone() - warp_vel = wp.to_torch(art_data.joint_vel).clone() + warp_pos = art_data.joint_pos.torch.clone() + warp_vel = art_data.joint_vel.torch.clone() # Run stable version (writes via write_joint_position_to_sim_index — which our mock # does not implement, so we compute the expected result directly) - defaults_t = wp.to_torch(art_data.default_joint_pos).clone() - limits_t = wp.to_torch(art_data.soft_joint_pos_limits) - vel_limits_t = wp.to_torch(art_data.soft_joint_vel_limits) + defaults_t = art_data.default_joint_pos.torch.clone() + limits_t = art_data.soft_joint_pos_limits.torch + vel_limits_t = art_data.soft_joint_vel_limits.torch expected_pos = defaults_t.clamp(limits_t[..., 0], limits_t[..., 1]) - expected_vel = wp.to_torch(art_data.default_joint_vel).clone().clamp(-vel_limits_t, vel_limits_t) + expected_vel = art_data.default_joint_vel.torch.clone().clamp(-vel_limits_t, vel_limits_t) assert_close(warp_pos, expected_pos) assert_close(warp_vel, expected_vel) @@ -185,15 +185,15 @@ def test_reset_joints_by_scale_parity(self, warp_env, stable_env, art_data, all_ warp_env, mask, position_range=(1.0, 1.0), velocity_range=(1.0, 1.0), asset_cfg=cfg ) wp.synchronize() - warp_pos = wp.to_torch(art_data.joint_pos).clone() - warp_vel = wp.to_torch(art_data.joint_vel).clone() + warp_pos = art_data.joint_pos.torch.clone() + warp_vel = art_data.joint_vel.torch.clone() # Expected: default * 1.0, clamped to limits - defaults_t = wp.to_torch(art_data.default_joint_pos).clone() - limits_t = wp.to_torch(art_data.soft_joint_pos_limits) - vel_limits_t = wp.to_torch(art_data.soft_joint_vel_limits) + defaults_t = art_data.default_joint_pos.torch.clone() + limits_t = art_data.soft_joint_pos_limits.torch + vel_limits_t = art_data.soft_joint_vel_limits.torch expected_pos = defaults_t.clamp(limits_t[..., 0], limits_t[..., 1]) - expected_vel = wp.to_torch(art_data.default_joint_vel).clone().clamp(-vel_limits_t, vel_limits_t) + expected_vel = art_data.default_joint_vel.torch.clone().clamp(-vel_limits_t, vel_limits_t) assert_close(warp_pos, expected_pos) assert_close(warp_vel, expected_vel) @@ -234,7 +234,7 @@ def test_reset_joints_by_offset(self, warp_env, art_data, all_joints_cfg): wp.synchronize() # With zero offset, joint_pos should equal new defaults (clamped to limits [-3.14, 3.14]) - result = wp.to_torch(art_data.joint_pos) + result = art_data.joint_pos.torch expected = torch.full((NUM_ENVS, NUM_JOINTS), 0.5, device=DEVICE) assert_close(result, expected) @@ -259,7 +259,7 @@ def test_reset_joints_by_scale(self, warp_env, art_data, all_joints_cfg): wp.capture_launch(cap.graph) wp.synchronize() - result = wp.to_torch(art_data.joint_pos) + result = art_data.joint_pos.torch expected = torch.full((NUM_ENVS, NUM_JOINTS), 0.25, device=DEVICE) assert_close(result, expected) @@ -333,7 +333,7 @@ def test_reset_joints_mask_selectivity(self, warp_env, art_data, all_joints_cfg) ) wp.synchronize() - result = wp.to_torch(art_data.joint_pos) + result = art_data.joint_pos.torch # Masked envs: reset to 0 (defaults + 0 offset) assert_close(result[: NUM_ENVS // 2], torch.zeros(NUM_ENVS // 2, NUM_JOINTS, device=DEVICE)) # Unmasked envs: still 999.0 diff --git a/source/isaaclab_mimic/isaaclab_mimic/datagen/data_generator.py b/source/isaaclab_mimic/isaaclab_mimic/datagen/data_generator.py index 6709ad2e6545..6851a90256f2 100644 --- a/source/isaaclab_mimic/isaaclab_mimic/datagen/data_generator.py +++ b/source/isaaclab_mimic/isaaclab_mimic/datagen/data_generator.py @@ -13,7 +13,6 @@ import numpy as np import torch -import warp as wp import isaaclab.utils.math as PoseUtils @@ -874,7 +873,7 @@ async def generate( # noqa: C901 # Update visualization if motion planner is available if motion_planner and motion_planner.visualize_spheres: - current_joints = wp.to_torch(self.env.scene["robot"].data.joint_pos)[env_id] + current_joints = self.env.scene["robot"].data.joint_pos.torch[env_id] motion_planner._update_visualization_at_joint_positions(current_joints) eef_waypoint_dict[eef_name] = waypoint diff --git a/source/isaaclab_mimic/isaaclab_mimic/locomanipulation_sdg/scene_utils.py b/source/isaaclab_mimic/isaaclab_mimic/locomanipulation_sdg/scene_utils.py index 2a05b5f70096..a0eb00fb4a58 100644 --- a/source/isaaclab_mimic/isaaclab_mimic/locomanipulation_sdg/scene_utils.py +++ b/source/isaaclab_mimic/isaaclab_mimic/locomanipulation_sdg/scene_utils.py @@ -85,7 +85,7 @@ def __init__(self, scene, entity_name: str, body_name: str): def get_pose(self): """Get the 3D pose of the entity.""" - body_link_state_w = wp.to_torch(self.scene[self.entity_name].data.body_link_state_w) + body_link_state_w = self.scene[self.entity_name].data.body_link_state_w.torch pose = body_link_state_w[ :, self.scene[self.entity_name].data.body_names.index(self.body_name), @@ -113,9 +113,9 @@ def _get_xform_view(self) -> FrameView: def get_pose(self): """Get the 3D pose of the entity.""" xform_prim = self._get_xform_view() - pos_wp, ori_wp = xform_prim.get_world_poses() - position = wp.to_torch(pos_wp) - orientation = wp.to_torch(ori_wp) + pos_w, quat_w = xform_prim.get_world_poses() + position = pos_w.torch + orientation = quat_w.torch pose = torch.cat([position, orientation], dim=-1) return pose diff --git a/source/isaaclab_mimic/isaaclab_mimic/motion_planners/curobo/curobo_planner.py b/source/isaaclab_mimic/isaaclab_mimic/motion_planners/curobo/curobo_planner.py index 3761a5858f26..c03db9c32921 100644 --- a/source/isaaclab_mimic/isaaclab_mimic/motion_planners/curobo/curobo_planner.py +++ b/source/isaaclab_mimic/isaaclab_mimic/motion_planners/curobo/curobo_planner.py @@ -9,7 +9,6 @@ import numpy as np import torch -import warp as wp from curobo.cuda_robot_model.cuda_robot_model import CudaRobotModelState from curobo.geom.sdf.world import CollisionCheckerType @@ -194,7 +193,7 @@ def __init__( # Use env-local base translation for multi-env rendering consistency env_origin = self.env.scene.env_origins[env_id, :3] - base_translation = (wp.to_torch(self.robot.data.root_pos_w)[env_id, :3] - env_origin).detach().cpu().numpy() + base_translation = (self.robot.data.root_pos_w.torch[env_id, :3] - env_origin).detach().cpu().numpy() self.plan_visualizer = PlanVisualizer( robot_name=self.config.robot_name, recording_id=f"curobo_plan_{env_id}", @@ -628,8 +627,8 @@ def _sync_object_poses_with_isaaclab(self) -> None: # Get current pose from Lab (may be on CPU or CUDA depending on --device flag) obj = rigid_objects[object_name] env_origin = self.env.scene.env_origins[self.env_id] - current_pos_raw = wp.to_torch(obj.data.root_pos_w)[self.env_id] - env_origin - current_quat_raw = wp.to_torch(obj.data.root_quat_w)[self.env_id] # (x, y, z, w) + current_pos_raw = obj.data.root_pos_w.torch[self.env_id] - env_origin + current_quat_raw = obj.data.root_quat_w.torch[self.env_id] # (x, y, z, w) # Convert to cuRobo device and extract float values for pose list current_pos = self._to_curobo_device(current_pos_raw) @@ -669,8 +668,8 @@ def _sync_object_poses_with_isaaclab(self) -> None: # Get current pose and update in collision checker obj = rigid_objects[object_name] env_origin = self.env.scene.env_origins[self.env_id] - current_pos_raw = wp.to_torch(obj.data.root_pos_w)[self.env_id] - env_origin - current_quat_raw = wp.to_torch(obj.data.root_quat_w)[self.env_id] + current_pos_raw = obj.data.root_pos_w.torch[self.env_id] - env_origin + current_quat_raw = obj.data.root_quat_w.torch[self.env_id] current_pos = self._to_curobo_device(current_pos_raw) current_quat = self._to_curobo_device(current_quat_raw) @@ -933,7 +932,7 @@ def _get_current_joint_state_for_curobo(self) -> JointState: and zero velocity/acceleration. """ # Fetch joint position (shape: [1, num_joints]) - joint_pos_raw: torch.Tensor = wp.to_torch(self.robot.data.joint_pos)[self.env_id, :].unsqueeze(0) + joint_pos_raw: torch.Tensor = self.robot.data.joint_pos.torch[self.env_id, :].unsqueeze(0) joint_vel_raw: torch.Tensor = torch.zeros_like(joint_pos_raw) joint_acc_raw: torch.Tensor = torch.zeros_like(joint_pos_raw) @@ -1574,7 +1573,7 @@ def _update_visualization_at_joint_positions(self, joint_positions: torch.Tensor if self.frame_counter % self.sphere_update_freq != 0: return - original_joints: torch.Tensor = wp.to_torch(self.robot.data.joint_pos)[self.env_id].clone() + original_joints: torch.Tensor = self.robot.data.joint_pos.torch[self.env_id].clone() try: # Ensure joint positions are on environment device for robot commands @@ -1701,7 +1700,7 @@ def _create_sphere_config(self, sphere_idx: int, sphere, robot_link_count: int) opacity = 0.9 if is_attached else 0.5 # Calculate position in world frame (do not use env_origin) - root_translation = wp.to_torch(self.robot.data.root_pos_w)[self.env_id, :3].detach().cpu().numpy() + root_translation = self.robot.data.root_pos_w.torch[self.env_id, :3].detach().cpu().numpy() position = sphere.position.cpu().numpy() if hasattr(sphere.position, "cpu") else sphere.position if not is_attached: position = position + root_translation @@ -1788,7 +1787,7 @@ def update_world_and_plan_motion( gripper_closed = expected_attached_object is not None self._set_gripper_state(gripper_closed) current_attached = self.get_attached_objects() - gripper_pos = wp.to_torch(self.robot.data.joint_pos)[env_id, -2:] + gripper_pos = self.robot.data.joint_pos.torch[env_id, -2:] self.logger.debug(f"Current attached objects: {current_attached}") @@ -1810,13 +1809,13 @@ def update_world_and_plan_motion( if expected_attached_object in rigid_objects: obj = rigid_objects[expected_attached_object] origin = self.env.scene.env_origins[env_id] - obj_pos = wp.to_torch(obj.data.root_pos_w)[env_id] - origin + obj_pos = obj.data.root_pos_w.torch[env_id] - origin self.logger.debug(f"Isaac Lab object position: {obj_pos}") # Debug end-effector position ee_frame_cfg = SceneEntityCfg("ee_frame") ee_frame = self.env.scene[ee_frame_cfg.name] - ee_pos = wp.to_torch(ee_frame.data.target_pos_w)[env_id, 0, :] - origin + ee_pos = ee_frame.data.target_pos_w.torch[env_id, 0, :] - origin self.logger.debug(f"End-effector position: {ee_pos}") # Debug distance diff --git a/source/isaaclab_mimic/test/test_curobo_planner_cube_stack.py b/source/isaaclab_mimic/test/test_curobo_planner_cube_stack.py index 04c5f4939f9c..8dea1cb14f74 100644 --- a/source/isaaclab_mimic/test/test_curobo_planner_cube_stack.py +++ b/source/isaaclab_mimic/test/test_curobo_planner_cube_stack.py @@ -27,7 +27,6 @@ import gymnasium as gym import torch -import warp as wp import isaaclab.utils.math as math_utils from isaaclab.assets import Articulation, RigidObject @@ -165,7 +164,7 @@ def _pose_from_xy_quat(self, xy: torch.Tensor, z: float, quat: torch.Tensor) -> def _get_cube_pos(self, cube_name: str) -> torch.Tensor: """Return the current world position of a cube's root (x, y, z).""" obj: RigidObject = self.env.scene[cube_name] - return wp.to_torch(obj.data.root_pos_w)[0, :3].clone().detach() + return obj.data.root_pos_w.torch[0, :3].clone().detach() def _place_pose_over_cube(self, cube_name: str, height_offset: float) -> torch.Tensor: """Compute a goal pose directly above the named cube using the latest pose.""" diff --git a/source/isaaclab_newton/config/extension.toml b/source/isaaclab_newton/config/extension.toml index 4329d2cd516f..e4bb2ccd0c74 100644 --- a/source/isaaclab_newton/config/extension.toml +++ b/source/isaaclab_newton/config/extension.toml @@ -1,7 +1,7 @@ [package] # Note: Semantic Versioning is used: https://semver.org/ -version = "0.5.22" +version = "0.5.24" # Description title = "Newton simulation interfaces for IsaacLab core package" diff --git a/source/isaaclab_newton/docs/CHANGELOG.rst b/source/isaaclab_newton/docs/CHANGELOG.rst index fe50e969092d..264a3c40cbd6 100644 --- a/source/isaaclab_newton/docs/CHANGELOG.rst +++ b/source/isaaclab_newton/docs/CHANGELOG.rst @@ -1,7 +1,7 @@ Changelog --------- -0.5.22 (2026-04-24) +0.5.24 (2026-04-27) ~~~~~~~~~~~~~~~~~~~ Added @@ -20,6 +20,37 @@ Added ``margin=0.01``. +0.5.23 (2026-04-24) +~~~~~~~~~~~~~~~~~~~ + +Changed +^^^^^^^ + +* Updated :class:`~isaaclab_newton.sim.views.NewtonSiteFrameView` to match the + new :class:`~isaaclab.sim.views.BaseFrameView` ProxyArray return contract. + See the ``isaaclab`` 4.6.15 changelog for migration guidance. + + +0.5.22 (2026-04-23) +~~~~~~~~~~~~~~~~~~~ + +Changed +^^^^^^^ + +* Properties on the following data classes now return + :class:`~isaaclab.utils.warp.ProxyArray` instead of raw ``wp.array``: + :class:`~isaaclab_newton.assets.articulation.ArticulationData`, + :class:`~isaaclab_newton.assets.rigid_object.RigidObjectData`, + :class:`~isaaclab_newton.assets.rigid_object_collection.RigidObjectCollectionData`, + :class:`~isaaclab_newton.sensors.contact_sensor.ContactSensorData`, + :class:`~isaaclab_newton.sensors.frame_transformer.FrameTransformerData`, + :class:`~isaaclab_newton.sensors.imu.ImuData`, and + :class:`~isaaclab_newton.sensors.pva.PvaData`. + Use ``.torch`` for a cached zero-copy ``torch.Tensor`` view, or ``.warp`` for + the underlying ``wp.array``. Implicit torch operations (arithmetic, + ``torch.*`` functions) work during the deprecation period but emit a warning. + + 0.5.21 (2026-04-23) ~~~~~~~~~~~~~~~~~~~ diff --git a/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py b/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py index 515176352490..39efda6cf7a9 100644 --- a/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py +++ b/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py @@ -3234,9 +3234,9 @@ def _create_buffers(self): # tendon names are set in _process_tendons function # -- joint commands (sent to the simulation after actuator processing) - self._joint_pos_target_sim = wp.zeros_like(self.data.joint_pos_target, device=self.device) - self._joint_vel_target_sim = wp.zeros_like(self.data.joint_pos_target, device=self.device) - self._joint_effort_target_sim = wp.zeros_like(self.data.joint_pos_target, device=self.device) + self._joint_pos_target_sim = wp.zeros_like(self.data.joint_pos_target.warp, device=self.device) + self._joint_vel_target_sim = wp.zeros_like(self.data.joint_pos_target.warp, device=self.device) + self._joint_effort_target_sim = wp.zeros_like(self.data.joint_pos_target.warp, device=self.device) # soft joint position limits (recommended not to be too close to limits). wp.launch( @@ -3340,12 +3340,12 @@ def _process_actuators_cfg(self): joint_ids=joint_ids, num_envs=self.num_instances, device=self.device, - stiffness=wp.to_torch(self._data.joint_stiffness)[:, joint_ids], - damping=wp.to_torch(self._data.joint_damping)[:, joint_ids], - armature=wp.to_torch(self._data.joint_armature)[:, joint_ids], - friction=wp.to_torch(self._data.joint_friction_coeff)[:, joint_ids], - effort_limit=wp.to_torch(self._data.joint_effort_limits)[:, joint_ids].clone(), - velocity_limit=wp.to_torch(self._data.joint_vel_limits)[:, joint_ids], + stiffness=self._data.joint_stiffness.torch[:, joint_ids], + damping=self._data.joint_damping.torch[:, joint_ids], + armature=self._data.joint_armature.torch[:, joint_ids], + friction=self._data.joint_friction_coeff.torch[:, joint_ids], + effort_limit=self._data.joint_effort_limits.torch[:, joint_ids].clone(), + velocity_limit=self._data.joint_vel_limits.torch[:, joint_ids], ) # store actuator group self.actuators[actuator_name] = actuator @@ -3472,16 +3472,16 @@ def _apply_actuator_model(self): # prepare input for actuator model based on cached data # TODO : A tensor dict would be nice to do the indexing of all tensors together control_action = ArticulationActions( - joint_positions=wp.to_torch(self._data.joint_pos_target)[:, actuator.joint_indices], - joint_velocities=wp.to_torch(self._data.joint_vel_target)[:, actuator.joint_indices], - joint_efforts=wp.to_torch(self._data.joint_effort_target)[:, actuator.joint_indices], + joint_positions=self._data.joint_pos_target.torch[:, actuator.joint_indices], + joint_velocities=self._data.joint_vel_target.torch[:, actuator.joint_indices], + joint_efforts=self._data.joint_effort_target.torch[:, actuator.joint_indices], joint_indices=actuator.joint_indices, ) # compute joint command from the actuator model control_action = actuator.compute( control_action, - joint_pos=wp.to_torch(self._data.joint_pos)[:, actuator.joint_indices], - joint_vel=wp.to_torch(self._data.joint_vel)[:, actuator.joint_indices], + joint_pos=self._data.joint_pos.torch[:, actuator.joint_indices], + joint_vel=self._data.joint_vel.torch[:, actuator.joint_indices], ) # update targets (these are set into the simulation) joint_indices = actuator.joint_indices @@ -3544,9 +3544,9 @@ def _validate_cfg(self): return # check that the default values are within the limits - joint_pos_limits_lower = wp.to_torch(self._data.joint_pos_limits_lower)[0] - joint_pos_limits_upper = wp.to_torch(self._data.joint_pos_limits_upper)[0] - default_joint_pos = wp.to_torch(self._data.default_joint_pos)[0] + joint_pos_limits_lower = self._data.joint_pos_limits_lower.torch[0] + joint_pos_limits_upper = self._data.joint_pos_limits_upper.torch[0] + default_joint_pos = self._data.default_joint_pos.torch[0] out_of_range = default_joint_pos < joint_pos_limits_lower out_of_range |= default_joint_pos > joint_pos_limits_upper violated_indices = torch.nonzero(out_of_range, as_tuple=False).squeeze(-1) @@ -3563,8 +3563,8 @@ def _validate_cfg(self): raise ValueError(msg) # check that the default joint velocities are within the limits - joint_max_vel = wp.to_torch(self._data.joint_vel_limits)[0] - default_joint_vel = wp.to_torch(self._data.default_joint_vel)[0] + joint_max_vel = self._data.joint_vel_limits.torch[0] + default_joint_vel = self._data.default_joint_vel.torch[0] out_of_range = torch.abs(default_joint_vel) > joint_max_vel violated_indices = torch.nonzero(out_of_range, as_tuple=False).squeeze(-1) if len(violated_indices) > 0: @@ -3604,19 +3604,19 @@ def format_limits(_, v: tuple[float, float]) -> str: # -- gains # Use data properties which have already been cloned and stored during initialization # This avoids issues with indexedarray or empty arrays from root_view - stiffnesses = wp.to_torch(self.data.joint_stiffness)[0].cpu().tolist() - dampings = wp.to_torch(self.data.joint_damping)[0].cpu().tolist() + stiffnesses = self.data.joint_stiffness.torch[0].cpu().tolist() + dampings = self.data.joint_damping.torch[0].cpu().tolist() # -- properties - armatures = wp.to_torch(self.data.joint_armature)[0].cpu().tolist() + armatures = self.data.joint_armature.torch[0].cpu().tolist() # For friction, use the individual components from data - friction_coeff = wp.to_torch(self.data.joint_friction_coeff)[0].cpu() + friction_coeff = self.data.joint_friction_coeff.torch[0].cpu() static_frictions = friction_coeff.tolist() # -- limits # joint_pos_limits is vec2f array, convert to torch and extract [lower, upper] pairs - position_limits_torch = wp.to_torch(self.data.joint_pos_limits)[0].cpu() # shape: (num_joints, 2) + position_limits_torch = self.data.joint_pos_limits.torch[0].cpu() # shape: (num_joints, 2) position_limits = [tuple(pos_limit.tolist()) for pos_limit in position_limits_torch] - velocity_limits = wp.to_torch(self.data.joint_vel_limits)[0].cpu().tolist() - effort_limits = wp.to_torch(self.data.joint_effort_limits)[0].cpu().tolist() + velocity_limits = self.data.joint_vel_limits.torch[0].cpu().tolist() + effort_limits = self.data.joint_effort_limits.torch[0].cpu().tolist() # create table for term information joint_table = PrettyTable() joint_table.title = f"Simulation Joint Information (Prim path: {self.cfg.prim_path})" diff --git a/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation_data.py b/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation_data.py index e78296c32864..ff74915db71b 100644 --- a/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation_data.py +++ b/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation_data.py @@ -15,6 +15,7 @@ from isaaclab.assets.articulation.base_articulation_data import BaseArticulationData from isaaclab.utils.buffers import TimestampedBufferWarp as TimestampedBuffer from isaaclab.utils.math import normalize +from isaaclab.utils.warp import ProxyArray from isaaclab.utils.warp.utils import capture_unsafe from isaaclab_newton.assets import kernels as shared_kernels @@ -82,8 +83,8 @@ def __init__(self, root_view: ArticulationView, device: str): forward_vec = torch.tensor((1.0, 0.0, 0.0), device=self.device).repeat(self._root_view.count, 1) # Initialize constants - self.GRAVITY_VEC_W = wp.from_torch(gravity_dir, dtype=wp.vec3f) - self.FORWARD_VEC_B = wp.from_torch(forward_vec, dtype=wp.vec3f) + self.GRAVITY_VEC_W = ProxyArray(wp.from_torch(gravity_dir, dtype=wp.vec3f)) + self.FORWARD_VEC_B = ProxyArray(wp.from_torch(forward_vec, dtype=wp.vec3f)) self._create_simulation_bindings() self._create_buffers() @@ -147,13 +148,13 @@ def update(self, dt: float) -> None: """ @property - def default_root_pose(self) -> wp.array: + def default_root_pose(self) -> ProxyArray: """Default root pose ``[pos, quat]`` in the local environment frame. The position and quaternion are of the articulation root's actor frame. Shape is (num_instances), dtype = wp.transformf. In torch this resolves to (num_instances, 7). """ - return self._default_root_pose + return self._default_root_pose_ta @default_root_pose.setter def default_root_pose(self, value: wp.array) -> None: @@ -170,13 +171,13 @@ def default_root_pose(self, value: wp.array) -> None: self._default_root_pose.assign(value) @property - def default_root_vel(self) -> wp.array: + def default_root_vel(self) -> ProxyArray: """Default root velocity ``[lin_vel, ang_vel]`` in the local environment frame. The linear and angular velocities are of the articulation root's center of mass frame. Shape is (num_instances), dtype = wp.spatial_vectorf. In torch this resolves to (num_instances, 6). """ - return self._default_root_vel + return self._default_root_vel_ta @default_root_vel.setter def default_root_vel(self, value: wp.array) -> None: @@ -193,14 +194,14 @@ def default_root_vel(self, value: wp.array) -> None: self._default_root_vel.assign(value) @property - def default_joint_pos(self) -> wp.array: + def default_joint_pos(self) -> ProxyArray: """Default joint positions of all joints. Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to (num_instances, num_joints). This quantity is configured through the :attr:`isaaclab.assets.ArticulationCfg.init_state` parameter. """ - return self._default_joint_pos + return self._default_joint_pos_ta @default_joint_pos.setter def default_joint_pos(self, value: wp.array) -> None: @@ -217,14 +218,14 @@ def default_joint_pos(self, value: wp.array) -> None: self._default_joint_pos.assign(value) @property - def default_joint_vel(self) -> wp.array: + def default_joint_vel(self) -> ProxyArray: """Default joint velocities of all joints. Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to (num_instances, num_joints). This quantity is configured through the :attr:`isaaclab.assets.ArticulationCfg.init_state` parameter. """ - return self._default_joint_vel + return self._default_joint_vel_ta @default_joint_vel.setter def default_joint_vel(self, value: wp.array) -> None: @@ -245,7 +246,7 @@ def default_joint_vel(self, value: wp.array) -> None: """ @property - def joint_pos_target(self) -> wp.array: + def joint_pos_target(self) -> ProxyArray: """Joint position targets commanded by the user. Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to (num_instances, num_joints). @@ -254,10 +255,10 @@ def joint_pos_target(self) -> wp.array: For an explicit actuator model, the targets are used to compute the joint torques (see :attr:`applied_torque`), which are then set into the simulation. """ - return self._joint_pos_target + return self._joint_pos_target_ta @property - def joint_vel_target(self) -> wp.array: + def joint_vel_target(self) -> ProxyArray: """Joint velocity targets commanded by the user. Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to (num_instances, num_joints). @@ -266,10 +267,10 @@ def joint_vel_target(self) -> wp.array: For an explicit actuator model, the targets are used to compute the joint torques (see :attr:`applied_torque`), which are then set into the simulation. """ - return self._joint_vel_target + return self._joint_vel_target_ta @property - def joint_effort_target(self) -> wp.array: + def joint_effort_target(self) -> ProxyArray: """Joint effort targets commanded by the user. Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to (num_instances, num_joints). @@ -278,14 +279,14 @@ def joint_effort_target(self) -> wp.array: For an explicit actuator model, the targets are used to compute the joint torques (see :attr:`applied_torque`), which are then set into the simulation. """ - return self._joint_effort_target + return self._joint_effort_target_ta """ Joint commands -- Explicit actuators. """ @property - def computed_torque(self) -> wp.array: + def computed_torque(self) -> ProxyArray: """Joint torques computed from the actuator model (before clipping). Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to (num_instances, num_joints). @@ -294,10 +295,10 @@ def computed_torque(self) -> wp.array: It is exposed for users who want to inspect the computations inside the actuator model. For instance, to penalize the learning agent for a difference between the computed and applied torques. """ - return self._computed_torque + return self._computed_torque_ta @property - def applied_torque(self) -> wp.array: + def applied_torque(self) -> ProxyArray: """Joint torques applied from the actuator model (after clipping). Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to (num_instances, num_joints). @@ -305,60 +306,60 @@ def applied_torque(self) -> wp.array: These torques are set into the simulation, after clipping the :attr:`computed_torque` based on the actuator model. """ - return self._applied_torque + return self._applied_torque_ta """ Joint properties """ @property - def joint_stiffness(self) -> wp.array: + def joint_stiffness(self) -> ProxyArray: """Joint stiffness provided to the simulation. Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to (num_instances, num_joints). In the case of explicit actuators, the value for the corresponding joints is zero. """ - return self._sim_bind_joint_stiffness_sim + return self._joint_stiffness_ta @property - def joint_damping(self) -> wp.array: + def joint_damping(self) -> ProxyArray: """Joint damping provided to the simulation. Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to (num_instances, num_joints). In the case of explicit actuators, the value for the corresponding joints is zero. """ - return self._sim_bind_joint_damping_sim + return self._joint_damping_ta @property - def joint_armature(self) -> wp.array: + def joint_armature(self) -> ProxyArray: """Joint armature provided to the simulation. Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to (num_instances, num_joints). """ - return self._sim_bind_joint_armature + return self._joint_armature_ta @property - def joint_friction_coeff(self) -> wp.array: + def joint_friction_coeff(self) -> ProxyArray: """Joint static friction coefficient provided to the simulation. Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to (num_instances, num_joints). """ - return self._sim_bind_joint_friction_coeff + return self._joint_friction_coeff_ta @property - def joint_pos_limits_lower(self) -> wp.array: + def joint_pos_limits_lower(self) -> ProxyArray: """Joint position limits lower provided to the simulation. Shape is (num_instances, num_joints).""" - return self._sim_bind_joint_pos_limits_lower + return self._joint_pos_limits_lower_ta @property - def joint_pos_limits_upper(self) -> wp.array: + def joint_pos_limits_upper(self) -> ProxyArray: """Joint position limits upper provided to the simulation. Shape is (num_instances, num_joints).""" - return self._sim_bind_joint_pos_limits_upper + return self._joint_pos_limits_upper_ta @property - def joint_pos_limits(self) -> wp.array: + def joint_pos_limits(self) -> ProxyArray: """Joint position limits provided to the simulation. Shape is (num_instances, num_joints, 2), dtype = wp.vec2f. In torch this resolves to @@ -370,6 +371,7 @@ def joint_pos_limits(self) -> wp.array: self._joint_pos_limits = wp.zeros( (self._num_instances, self._num_joints), dtype=wp.vec2f, device=self.device ) + self._joint_pos_limits_ta = ProxyArray(self._joint_pos_limits) wp.launch( articulation_kernels.concat_joint_pos_limits_lower_and_upper, dim=(self._num_instances, self._num_joints), @@ -382,30 +384,30 @@ def joint_pos_limits(self) -> wp.array: ], device=self.device, ) - return self._joint_pos_limits + return self._joint_pos_limits_ta @property - def joint_vel_limits(self) -> wp.array: + def joint_vel_limits(self) -> ProxyArray: """Joint maximum velocity provided to the simulation. Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to (num_instances, num_joints). """ - return self._sim_bind_joint_vel_limits_sim + return self._joint_vel_limits_ta @property - def joint_effort_limits(self) -> wp.array: + def joint_effort_limits(self) -> ProxyArray: """Joint maximum effort provided to the simulation. Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to (num_instances, num_joints). """ - return self._sim_bind_joint_effort_limits_sim + return self._joint_effort_limits_ta """ Joint properties - Custom. """ @property - def soft_joint_pos_limits(self) -> wp.array: + def soft_joint_pos_limits(self) -> ProxyArray: r"""Soft joint positions limits for all joints. Shape is (num_instances, num_joints), dtype = wp.vec2f. In torch this resolves to @@ -426,10 +428,10 @@ def soft_joint_pos_limits(self) -> wp.array: The soft joint position limits help specify a safety region around the joint limits. It isn't used by the simulation, but is useful for learning agents to prevent the joint positions from violating the limits. """ - return self._soft_joint_pos_limits + return self._soft_joint_pos_limits_ta @property - def soft_joint_vel_limits(self) -> wp.array: + def soft_joint_vel_limits(self) -> ProxyArray: """Soft joint velocity limits for all joints. Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to (num_instances, num_joints). @@ -437,22 +439,22 @@ def soft_joint_vel_limits(self) -> wp.array: These are obtained from the actuator model. It may differ from :attr:`joint_vel_limits` if the actuator model has a variable velocity limit model. For instance, in a variable gear ratio actuator model. """ - return self._soft_joint_vel_limits + return self._soft_joint_vel_limits_ta @property - def gear_ratio(self) -> wp.array: + def gear_ratio(self) -> ProxyArray: """Gear ratio for relating motor torques to applied Joint torques. Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to (num_instances, num_joints). """ - return self._gear_ratio + return self._gear_ratio_ta """ Fixed tendon properties. """ @property - def fixed_tendon_stiffness(self) -> wp.array: + def fixed_tendon_stiffness(self) -> ProxyArray: """Fixed tendon stiffness provided to the simulation. Shape is (num_instances, num_fixed_tendons), dtype = wp.float32. In torch this resolves to @@ -461,7 +463,7 @@ def fixed_tendon_stiffness(self) -> wp.array: raise NotImplementedError @property - def fixed_tendon_damping(self) -> wp.array: + def fixed_tendon_damping(self) -> ProxyArray: """Fixed tendon damping provided to the simulation. Shape is (num_instances, num_fixed_tendons), dtype = wp.float32. In torch this resolves to @@ -470,7 +472,7 @@ def fixed_tendon_damping(self) -> wp.array: raise NotImplementedError @property - def fixed_tendon_limit_stiffness(self) -> wp.array: + def fixed_tendon_limit_stiffness(self) -> ProxyArray: """Fixed tendon limit stiffness provided to the simulation. Shape is (num_instances, num_fixed_tendons), dtype = wp.float32. In torch this resolves to @@ -479,7 +481,7 @@ def fixed_tendon_limit_stiffness(self) -> wp.array: raise NotImplementedError @property - def fixed_tendon_rest_length(self) -> wp.array: + def fixed_tendon_rest_length(self) -> ProxyArray: """Fixed tendon rest length provided to the simulation. Shape is (num_instances, num_fixed_tendons), dtype = wp.float32. In torch this resolves to @@ -488,7 +490,7 @@ def fixed_tendon_rest_length(self) -> wp.array: raise NotImplementedError @property - def fixed_tendon_offset(self) -> wp.array: + def fixed_tendon_offset(self) -> ProxyArray: """Fixed tendon offset provided to the simulation. Shape is (num_instances, num_fixed_tendons), dtype = wp.float32. In torch this resolves to @@ -497,7 +499,7 @@ def fixed_tendon_offset(self) -> wp.array: raise NotImplementedError @property - def fixed_tendon_pos_limits(self) -> wp.array: + def fixed_tendon_pos_limits(self) -> ProxyArray: """Fixed tendon position limits provided to the simulation. Shape is (num_instances, num_fixed_tendons, 2), dtype = wp.vec2f. In torch this resolves to @@ -510,7 +512,7 @@ def fixed_tendon_pos_limits(self) -> wp.array: """ @property - def spatial_tendon_stiffness(self) -> wp.array: + def spatial_tendon_stiffness(self) -> ProxyArray: """Spatial tendon stiffness provided to the simulation. Shape is (num_instances, num_spatial_tendons), dtype = wp.float32. In torch this resolves to @@ -519,7 +521,7 @@ def spatial_tendon_stiffness(self) -> wp.array: raise NotImplementedError @property - def spatial_tendon_damping(self) -> wp.array: + def spatial_tendon_damping(self) -> ProxyArray: """Spatial tendon damping provided to the simulation. Shape is (num_instances, num_spatial_tendons), dtype = wp.float32. In torch this resolves to @@ -528,7 +530,7 @@ def spatial_tendon_damping(self) -> wp.array: raise NotImplementedError @property - def spatial_tendon_limit_stiffness(self) -> wp.array: + def spatial_tendon_limit_stiffness(self) -> ProxyArray: """Spatial tendon limit stiffness provided to the simulation. Shape is (num_instances, num_spatial_tendons), dtype = wp.float32. In torch this resolves to @@ -537,7 +539,7 @@ def spatial_tendon_limit_stiffness(self) -> wp.array: raise NotImplementedError @property - def spatial_tendon_offset(self) -> wp.array: + def spatial_tendon_offset(self) -> ProxyArray: """Spatial tendon offset provided to the simulation. Shape is (num_instances, num_spatial_tendons), dtype = wp.float32. In torch this resolves to @@ -550,7 +552,7 @@ def spatial_tendon_offset(self) -> wp.array: """ @property - def root_link_pose_w(self) -> wp.array: + def root_link_pose_w(self) -> ProxyArray: """Root link pose ``[pos, quat]`` in simulation world frame. Shape is (num_instances,), dtype = wp.transformf. In torch this resolves to (num_instances, 7). @@ -558,11 +560,11 @@ def root_link_pose_w(self) -> wp.array: This quantity is the pose of the articulation root's actor frame relative to the world. The orientation is provided in (x, y, z, w) format. """ - return self._sim_bind_root_link_pose_w + return self._root_link_pose_w_ta @property @capture_unsafe(_LAZY_CAPTURE_REASON) - def root_link_vel_w(self) -> wp.array: + def root_link_vel_w(self) -> ProxyArray: """Root link velocity ``[lin_vel, ang_vel]`` in simulation world frame. Shape is (num_instances,), dtype = wp.spatial_vectorf. In torch this resolves to (num_instances, 6). @@ -575,9 +577,9 @@ def root_link_vel_w(self) -> wp.array: shared_kernels.get_root_link_vel_from_root_com_vel, dim=self._num_instances, inputs=[ - self.root_com_vel_w, - self.root_link_pose_w, - self.body_com_pos_b, + self.root_com_vel_w.warp, + self.root_link_pose_w.warp, + self.body_com_pos_b.warp, ], outputs=[ self._root_link_vel_w.data, @@ -586,11 +588,11 @@ def root_link_vel_w(self) -> wp.array: ) self._root_link_vel_w.timestamp = self._sim_timestamp - return self._root_link_vel_w.data + return self._root_link_vel_w_ta @property @capture_unsafe(_LAZY_CAPTURE_REASON) - def root_com_pose_w(self) -> wp.array: + def root_com_pose_w(self) -> ProxyArray: """Root center of mass pose ``[pos, quat]`` in simulation world frame. Shape is (num_instances,), dtype = wp.transformf. In torch this resolves to (num_instances, 7). @@ -604,8 +606,8 @@ def root_com_pose_w(self) -> wp.array: shared_kernels.get_root_com_pose_from_root_link_pose, dim=self._num_instances, inputs=[ - self.root_link_pose_w, - self.body_com_pos_b, + self.root_link_pose_w.warp, + self.body_com_pos_b.warp, ], outputs=[ self._root_com_pose_w.data, @@ -614,10 +616,10 @@ def root_com_pose_w(self) -> wp.array: ) self._root_com_pose_w.timestamp = self._sim_timestamp - return self._root_com_pose_w.data + return self._root_com_pose_w_ta @property - def root_com_vel_w(self) -> wp.array: + def root_com_vel_w(self) -> ProxyArray: """Root center of mass velocity ``[lin_vel, ang_vel]`` in simulation world frame. Shape is (num_instances,), dtype = wp.spatial_vectorf. In torch this resolves to (num_instances, 6). @@ -625,31 +627,31 @@ def root_com_vel_w(self) -> wp.array: This quantity contains the linear and angular velocities of the articulation root's center of mass frame relative to the world. """ - return self._sim_bind_root_com_vel_w + return self._root_com_vel_w_ta """ Body state properties. """ @property - def body_mass(self) -> wp.array: + def body_mass(self) -> ProxyArray: """Body mass ``wp.float32`` in the world frame. Shape is (num_instances, num_bodies), dtype = wp.float32. In torch this resolves to (num_instances, num_bodies). """ - return self._sim_bind_body_mass + return self._body_mass_ta @property - def body_inertia(self) -> wp.array: + def body_inertia(self) -> ProxyArray: """Flattened body inertia in the world frame. Shape is (num_instances, num_bodies, 9), dtype = wp.float32. In torch this resolves to (num_instances, num_bodies, 9). """ - return self._sim_bind_body_inertia + return self._body_inertia_ta @property - def body_link_pose_w(self) -> wp.array: + def body_link_pose_w(self) -> ProxyArray: """Body link pose ``[pos, quat]`` in simulation world frame. Shape is (num_instances, num_bodies), dtype = wp.transformf. In torch this resolves to @@ -661,11 +663,11 @@ def body_link_pose_w(self) -> wp.array: if self._fk_timestamp < self._sim_timestamp: SimulationManager.forward() self._fk_timestamp = self._sim_timestamp - return self._sim_bind_body_link_pose_w + return self._body_link_pose_w_ta @property @capture_unsafe(_LAZY_CAPTURE_REASON) - def body_link_vel_w(self) -> wp.array: + def body_link_vel_w(self) -> ProxyArray: """Body link velocity ``[lin_vel, ang_vel]`` in simulation world frame. Shape is (num_instances, num_bodies), dtype = wp.spatial_vectorf. In torch this resolves to @@ -679,9 +681,9 @@ def body_link_vel_w(self) -> wp.array: shared_kernels.get_body_link_vel_from_body_com_vel, dim=(self._num_instances, self._num_bodies), inputs=[ - self.body_com_vel_w, - self.body_link_pose_w, - self.body_com_pos_b, + self.body_com_vel_w.warp, + self.body_link_pose_w.warp, + self.body_com_pos_b.warp, ], outputs=[ self._body_link_vel_w.data, @@ -690,11 +692,11 @@ def body_link_vel_w(self) -> wp.array: ) self._body_link_vel_w.timestamp = self._sim_timestamp - return self._body_link_vel_w.data + return self._body_link_vel_w_ta @property @capture_unsafe(_LAZY_CAPTURE_REASON) - def body_com_pose_w(self) -> wp.array: + def body_com_pose_w(self) -> ProxyArray: """Body center of mass pose ``[pos, quat]`` in simulation world frame. Shape is (num_instances, num_bodies), dtype = wp.transformf. In torch this resolves to @@ -708,8 +710,8 @@ def body_com_pose_w(self) -> wp.array: shared_kernels.get_body_com_pose_from_body_link_pose, dim=(self._num_instances, self._num_bodies), inputs=[ - self.body_link_pose_w, - self.body_com_pos_b, + self.body_link_pose_w.warp, + self.body_com_pos_b.warp, ], outputs=[ self._body_com_pose_w.data, @@ -718,10 +720,10 @@ def body_com_pose_w(self) -> wp.array: ) self._body_com_pose_w.timestamp = self._sim_timestamp - return self._body_com_pose_w.data + return self._body_com_pose_w_ta @property - def body_com_vel_w(self) -> wp.array: + def body_com_vel_w(self) -> ProxyArray: """Body center of mass velocity ``[lin_vel, ang_vel]`` in simulation world frame. Shape is (num_instances, num_bodies), dtype = wp.spatial_vectorf. In torch this resolves to @@ -730,10 +732,10 @@ def body_com_vel_w(self) -> wp.array: This quantity contains the linear and angular velocities of the articulation links' center of mass frame relative to the world. """ - return self._sim_bind_body_com_vel_w + return self._body_com_vel_w_ta @property - def body_com_acc_w(self): + def body_com_acc_w(self) -> ProxyArray: """Acceleration of all bodies center of mass ``[lin_acc, ang_acc]``. Shape is (num_instances, num_bodies), dtype = wp.spatial_vectorf. In torch this resolves to @@ -758,10 +760,10 @@ def body_com_acc_w(self): # set the buffer data and timestamp self._body_com_acc_w.timestamp = self._sim_timestamp # update the previous velocity - return self._body_com_acc_w.data + return self._body_com_acc_w_ta @property - def body_com_pos_b(self) -> wp.array: + def body_com_pos_b(self) -> ProxyArray: """Center of mass position of all of the bodies in their respective link frames. Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to @@ -769,10 +771,10 @@ def body_com_pos_b(self) -> wp.array: This quantity is the center of mass location relative to its body's link frame. """ - return self._sim_bind_body_com_pos_b + return self._body_com_pos_b_ta @property - def body_com_pose_b(self) -> wp.array: + def body_com_pose_b(self) -> ProxyArray: """Center of mass pose ``[pos, quat]`` of all bodies in their respective body's link frames. Shape is (num_instances, num_bodies), dtype = wp.transformf. In torch this resolves to @@ -793,7 +795,7 @@ def body_com_pose_b(self) -> wp.array: shared_kernels.make_dummy_body_com_pose_b, dim=(self._num_instances, self._num_bodies), inputs=[ - self.body_com_pos_b, + self.body_com_pos_b.warp, ], outputs=[ self._body_com_pose_b.data, @@ -801,10 +803,10 @@ def body_com_pose_b(self) -> wp.array: device=self.device, ) self._body_com_pose_b.timestamp = self._sim_timestamp - return self._body_com_pose_b.data + return self._body_com_pose_b_ta @property - def body_incoming_joint_wrench_b(self) -> wp.array: + def body_incoming_joint_wrench_b(self) -> ProxyArray: """Joint reaction wrench applied from body parent to child body in parent body frame. Shape is (num_instances, num_bodies), dtype = wp.spatial_vectorf. In torch this resolves to @@ -824,25 +826,25 @@ def body_incoming_joint_wrench_b(self) -> wp.array: """ @property - def joint_pos(self) -> wp.array: + def joint_pos(self) -> ProxyArray: """Joint positions of all joints. Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to (num_instances, num_joints). """ - return self._sim_bind_joint_pos + return self._joint_pos_ta @property - def joint_vel(self) -> wp.array: + def joint_vel(self) -> ProxyArray: """Joint velocities of all joints. Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to (num_instances, num_joints). """ - return self._sim_bind_joint_vel + return self._joint_vel_ta @property - def joint_acc(self) -> wp.array: + def joint_acc(self) -> ProxyArray: """Joint acceleration of all joints. Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to @@ -855,7 +857,7 @@ def joint_acc(self) -> wp.array: articulation_kernels.get_joint_acc_from_joint_vel, dim=(self._num_instances, self._num_joints), inputs=[ - self.joint_vel, + self.joint_vel.warp, self._previous_joint_vel, time_elapsed, ], @@ -865,7 +867,7 @@ def joint_acc(self) -> wp.array: device=self.device, ) self._joint_acc.timestamp = self._sim_timestamp - return self._joint_acc.data + return self._joint_acc_ta """ Derived Properties. @@ -873,7 +875,7 @@ def joint_acc(self) -> wp.array: @property @capture_unsafe(_LAZY_CAPTURE_REASON) - def projected_gravity_b(self): + def projected_gravity_b(self) -> ProxyArray: """Projection of the gravity direction on base frame. Shape is (num_instances), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). @@ -882,16 +884,16 @@ def projected_gravity_b(self): wp.launch( shared_kernels.quat_apply_inverse_1D_kernel, dim=self._num_instances, - inputs=[self.GRAVITY_VEC_W, self.root_link_quat_w], + inputs=[self.GRAVITY_VEC_W.warp, self.root_link_quat_w.warp], outputs=[self._projected_gravity_b.data], device=self.device, ) self._projected_gravity_b.timestamp = self._sim_timestamp - return self._projected_gravity_b.data + return self._projected_gravity_b_ta @property @capture_unsafe(_LAZY_CAPTURE_REASON) - def heading_w(self): + def heading_w(self) -> ProxyArray: """Yaw heading of the base frame (in radians). Shape is (num_instances), dtype = wp.float32. In torch this resolves to (num_instances,). @@ -904,16 +906,16 @@ def heading_w(self): wp.launch( shared_kernels.root_heading_w, dim=self._num_instances, - inputs=[self.FORWARD_VEC_B, self.root_link_quat_w], + inputs=[self.FORWARD_VEC_B.warp, self.root_link_quat_w.warp], outputs=[self._heading_w.data], device=self.device, ) self._heading_w.timestamp = self._sim_timestamp - return self._heading_w.data + return self._heading_w_ta @property @capture_unsafe(_LAZY_CAPTURE_REASON) - def root_link_lin_vel_b(self) -> wp.array: + def root_link_lin_vel_b(self) -> ProxyArray: """Root link linear velocity in base frame. Shape is (num_instances), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). @@ -925,19 +927,20 @@ def root_link_lin_vel_b(self) -> wp.array: self._root_link_lin_vel_b = TimestampedBuffer( shape=(self._num_instances,), dtype=wp.vec3f, device=self.device ) + self._root_link_lin_vel_b_ta = ProxyArray(self._root_link_lin_vel_b.data) if self._root_link_lin_vel_b.timestamp < self._sim_timestamp: wp.launch( shared_kernels.quat_apply_inverse_1D_kernel, dim=self._num_instances, - inputs=[self.root_link_lin_vel_w, self.root_link_quat_w], + inputs=[self.root_link_lin_vel_w.warp, self.root_link_quat_w.warp], outputs=[self._root_link_lin_vel_b.data], device=self.device, ) self._root_link_lin_vel_b.timestamp = self._sim_timestamp - return self._root_link_lin_vel_b.data + return self._root_link_lin_vel_b_ta @property - def root_link_ang_vel_b(self) -> wp.array: + def root_link_ang_vel_b(self) -> ProxyArray: """Root link angular velocity in base frame. Shape is (num_instances), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). @@ -949,19 +952,20 @@ def root_link_ang_vel_b(self) -> wp.array: self._root_link_ang_vel_b = TimestampedBuffer( shape=(self._num_instances,), dtype=wp.vec3f, device=self.device ) + self._root_link_ang_vel_b_ta = ProxyArray(self._root_link_ang_vel_b.data) if self._root_link_ang_vel_b.timestamp < self._sim_timestamp: wp.launch( shared_kernels.quat_apply_inverse_1D_kernel, dim=self._num_instances, - inputs=[self.root_link_ang_vel_w, self.root_link_quat_w], + inputs=[self.root_link_ang_vel_w.warp, self.root_link_quat_w.warp], outputs=[self._root_link_ang_vel_b.data], device=self.device, ) self._root_link_ang_vel_b.timestamp = self._sim_timestamp - return self._root_link_ang_vel_b.data + return self._root_link_ang_vel_b_ta @property - def root_com_lin_vel_b(self) -> wp.array: + def root_com_lin_vel_b(self) -> ProxyArray: """Root center of mass linear velocity in base frame. Shape is (num_instances), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). @@ -973,19 +977,20 @@ def root_com_lin_vel_b(self) -> wp.array: self._root_com_lin_vel_b = TimestampedBuffer( shape=(self._num_instances,), dtype=wp.vec3f, device=self.device ) + self._root_com_lin_vel_b_ta = ProxyArray(self._root_com_lin_vel_b.data) if self._root_com_lin_vel_b.timestamp < self._sim_timestamp: wp.launch( shared_kernels.quat_apply_inverse_1D_kernel, dim=self._num_instances, - inputs=[self.root_com_lin_vel_w, self.root_link_quat_w], + inputs=[self.root_com_lin_vel_w.warp, self.root_link_quat_w.warp], outputs=[self._root_com_lin_vel_b.data], device=self.device, ) self._root_com_lin_vel_b.timestamp = self._sim_timestamp - return self._root_com_lin_vel_b.data + return self._root_com_lin_vel_b_ta @property - def root_com_ang_vel_b(self) -> wp.array: + def root_com_ang_vel_b(self) -> ProxyArray: """Root center of mass angular velocity in base frame. Shape is (num_instances), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). @@ -997,103 +1002,134 @@ def root_com_ang_vel_b(self) -> wp.array: self._root_com_ang_vel_b = TimestampedBuffer( shape=(self._num_instances,), dtype=wp.vec3f, device=self.device ) + self._root_com_ang_vel_b_ta = ProxyArray(self._root_com_ang_vel_b.data) if self._root_com_ang_vel_b.timestamp < self._sim_timestamp: wp.launch( shared_kernels.quat_apply_inverse_1D_kernel, dim=self._num_instances, - inputs=[self.root_com_ang_vel_w, self.root_link_quat_w], + inputs=[self.root_com_ang_vel_w.warp, self.root_link_quat_w.warp], outputs=[self._root_com_ang_vel_b.data], device=self.device, ) self._root_com_ang_vel_b.timestamp = self._sim_timestamp - return self._root_com_ang_vel_b.data + return self._root_com_ang_vel_b_ta """ Sliced properties. """ @property - def root_link_pos_w(self) -> wp.array: + def root_link_pos_w(self) -> ProxyArray: """Root link position in simulation world frame. Shape is (num_instances), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). This quantity is the position of the actor frame of the root rigid body relative to the world. """ - return self._get_pos_from_transform(self._root_link_pos_w, self.root_link_pose_w) + self._root_link_pos_w = self._get_pos_from_transform(self._root_link_pos_w, self.root_link_pose_w.warp) + if self._root_link_pos_w_ta is None: + self._root_link_pos_w_ta = ProxyArray(self._root_link_pos_w) + return self._root_link_pos_w_ta @property - def root_link_quat_w(self) -> wp.array: + def root_link_quat_w(self) -> ProxyArray: """Root link orientation (x, y, z, w) in simulation world frame. Shape is (num_instances), dtype = wp.quatf. In torch this resolves to (num_instances, 4). This quantity is the orientation of the actor frame of the root rigid body. """ - return self._get_quat_from_transform(self._root_link_quat_w, self.root_link_pose_w) + self._root_link_quat_w = self._get_quat_from_transform(self._root_link_quat_w, self.root_link_pose_w.warp) + if self._root_link_quat_w_ta is None: + self._root_link_quat_w_ta = ProxyArray(self._root_link_quat_w) + return self._root_link_quat_w_ta @property - def root_link_lin_vel_w(self) -> wp.array: + def root_link_lin_vel_w(self) -> ProxyArray: """Root linear velocity in simulation world frame. Shape is (num_instances), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). This quantity is the linear velocity of the root rigid body's actor frame relative to the world. """ - return self._get_top_from_spatial_vector(self._root_link_lin_vel_w, self.root_link_vel_w) + self._root_link_lin_vel_w = self._get_top_from_spatial_vector( + self._root_link_lin_vel_w, self.root_link_vel_w.warp + ) + if self._root_link_lin_vel_w_ta is None: + self._root_link_lin_vel_w_ta = ProxyArray(self._root_link_lin_vel_w) + return self._root_link_lin_vel_w_ta @property - def root_link_ang_vel_w(self) -> wp.array: + def root_link_ang_vel_w(self) -> ProxyArray: """Root link angular velocity in simulation world frame. Shape is (num_instances), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). This quantity is the angular velocity of the actor frame of the root rigid body relative to the world. """ - return self._get_bottom_from_spatial_vector(self._root_link_ang_vel_w, self.root_link_vel_w) + self._root_link_ang_vel_w = self._get_bottom_from_spatial_vector( + self._root_link_ang_vel_w, self.root_link_vel_w.warp + ) + if self._root_link_ang_vel_w_ta is None: + self._root_link_ang_vel_w_ta = ProxyArray(self._root_link_ang_vel_w) + return self._root_link_ang_vel_w_ta @property - def root_com_pos_w(self) -> wp.array: + def root_com_pos_w(self) -> ProxyArray: """Root center of mass position in simulation world frame. Shape is (num_instances), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). This quantity is the position of the center of mass frame of the root rigid body relative to the world. """ - return self._get_pos_from_transform(self._root_com_pos_w, self.root_com_pose_w) + self._root_com_pos_w = self._get_pos_from_transform(self._root_com_pos_w, self.root_com_pose_w.warp) + if self._root_com_pos_w_ta is None: + self._root_com_pos_w_ta = ProxyArray(self._root_com_pos_w) + return self._root_com_pos_w_ta @property - def root_com_quat_w(self) -> wp.array: + def root_com_quat_w(self) -> ProxyArray: """Root center of mass orientation (x, y, z, w) in simulation world frame. Shape is (num_instances), dtype = wp.quatf. In torch this resolves to (num_instances, 4). This quantity is the orientation of the principal axes of inertia of the root rigid body relative to the world. """ - return self._get_quat_from_transform(self._root_com_quat_w, self.root_com_pose_w) + self._root_com_quat_w = self._get_quat_from_transform(self._root_com_quat_w, self.root_com_pose_w.warp) + if self._root_com_quat_w_ta is None: + self._root_com_quat_w_ta = ProxyArray(self._root_com_quat_w) + return self._root_com_quat_w_ta @property - def root_com_lin_vel_w(self) -> wp.array: + def root_com_lin_vel_w(self) -> ProxyArray: """Root center of mass linear velocity in simulation world frame. Shape is (num_instances), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). This quantity is the linear velocity of the root rigid body's center of mass frame relative to the world. """ - return self._get_top_from_spatial_vector(self._root_com_lin_vel_w, self.root_com_vel_w) + self._root_com_lin_vel_w = self._get_top_from_spatial_vector(self._root_com_lin_vel_w, self.root_com_vel_w.warp) + if self._root_com_lin_vel_w_ta is None: + self._root_com_lin_vel_w_ta = ProxyArray(self._root_com_lin_vel_w) + return self._root_com_lin_vel_w_ta @property - def root_com_ang_vel_w(self) -> wp.array: + def root_com_ang_vel_w(self) -> ProxyArray: """Root center of mass angular velocity in simulation world frame. Shape is (num_instances), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). This quantity is the angular velocity of the root rigid body's center of mass frame relative to the world. """ - return self._get_bottom_from_spatial_vector(self._root_com_ang_vel_w, self.root_com_vel_w) + self._root_com_ang_vel_w = self._get_bottom_from_spatial_vector( + self._root_com_ang_vel_w, self.root_com_vel_w.warp + ) + if self._root_com_ang_vel_w_ta is None: + self._root_com_ang_vel_w_ta = ProxyArray(self._root_com_ang_vel_w) + return self._root_com_ang_vel_w_ta @property - def body_link_pos_w(self) -> wp.array: + def body_link_pos_w(self) -> ProxyArray: """Positions of all bodies in simulation world frame. Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to @@ -1101,10 +1137,13 @@ def body_link_pos_w(self) -> wp.array: This quantity is the position of the articulation bodies' actor frame relative to the world. """ - return self._get_pos_from_transform(self._body_link_pos_w, self.body_link_pose_w) + self._body_link_pos_w = self._get_pos_from_transform(self._body_link_pos_w, self.body_link_pose_w.warp) + if self._body_link_pos_w_ta is None: + self._body_link_pos_w_ta = ProxyArray(self._body_link_pos_w) + return self._body_link_pos_w_ta @property - def body_link_quat_w(self) -> wp.array: + def body_link_quat_w(self) -> ProxyArray: """Orientation (x, y, z, w) of all bodies in simulation world frame. Shape is (num_instances, num_bodies), dtype = wp.quatf. In torch this resolves to @@ -1112,10 +1151,13 @@ def body_link_quat_w(self) -> wp.array: This quantity is the orientation of the articulation bodies' actor frame relative to the world. """ - return self._get_quat_from_transform(self._body_link_quat_w, self.body_link_pose_w) + self._body_link_quat_w = self._get_quat_from_transform(self._body_link_quat_w, self.body_link_pose_w.warp) + if self._body_link_quat_w_ta is None: + self._body_link_quat_w_ta = ProxyArray(self._body_link_quat_w) + return self._body_link_quat_w_ta @property - def body_link_lin_vel_w(self) -> wp.array: + def body_link_lin_vel_w(self) -> ProxyArray: """Linear velocity of all bodies in simulation world frame. Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to @@ -1123,10 +1165,15 @@ def body_link_lin_vel_w(self) -> wp.array: This quantity is the linear velocity of the articulation bodies' actor frame relative to the world. """ - return self._get_top_from_spatial_vector(self._body_link_lin_vel_w, self.body_link_vel_w) + self._body_link_lin_vel_w = self._get_top_from_spatial_vector( + self._body_link_lin_vel_w, self.body_link_vel_w.warp + ) + if self._body_link_lin_vel_w_ta is None: + self._body_link_lin_vel_w_ta = ProxyArray(self._body_link_lin_vel_w) + return self._body_link_lin_vel_w_ta @property - def body_link_ang_vel_w(self) -> wp.array: + def body_link_ang_vel_w(self) -> ProxyArray: """Angular velocity of all bodies in simulation world frame. Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to @@ -1134,10 +1181,15 @@ def body_link_ang_vel_w(self) -> wp.array: This quantity is the angular velocity of the articulation bodies' actor frame relative to the world. """ - return self._get_bottom_from_spatial_vector(self._body_link_ang_vel_w, self.body_link_vel_w) + self._body_link_ang_vel_w = self._get_bottom_from_spatial_vector( + self._body_link_ang_vel_w, self.body_link_vel_w.warp + ) + if self._body_link_ang_vel_w_ta is None: + self._body_link_ang_vel_w_ta = ProxyArray(self._body_link_ang_vel_w) + return self._body_link_ang_vel_w_ta @property - def body_com_pos_w(self) -> wp.array: + def body_com_pos_w(self) -> ProxyArray: """Positions of all bodies in simulation world frame. Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to @@ -1145,10 +1197,13 @@ def body_com_pos_w(self) -> wp.array: This quantity is the position of the articulation bodies' center of mass frame. """ - return self._get_pos_from_transform(self._body_com_pos_w, self.body_com_pose_w) + self._body_com_pos_w = self._get_pos_from_transform(self._body_com_pos_w, self.body_com_pose_w.warp) + if self._body_com_pos_w_ta is None: + self._body_com_pos_w_ta = ProxyArray(self._body_com_pos_w) + return self._body_com_pos_w_ta @property - def body_com_quat_w(self) -> wp.array: + def body_com_quat_w(self) -> ProxyArray: """Orientation (x, y, z, w) of the principal axes of inertia of all bodies in simulation world frame. Shape is (num_instances, num_bodies), dtype = wp.quatf. In torch this resolves to @@ -1156,10 +1211,13 @@ def body_com_quat_w(self) -> wp.array: This quantity is the orientation of the principal axes of inertia of the articulation bodies. """ - return self._get_quat_from_transform(self._body_com_quat_w, self.body_com_pose_w) + self._body_com_quat_w = self._get_quat_from_transform(self._body_com_quat_w, self.body_com_pose_w.warp) + if self._body_com_quat_w_ta is None: + self._body_com_quat_w_ta = ProxyArray(self._body_com_quat_w) + return self._body_com_quat_w_ta @property - def body_com_lin_vel_w(self) -> wp.array: + def body_com_lin_vel_w(self) -> ProxyArray: """Linear velocity of all bodies in simulation world frame. Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to @@ -1167,10 +1225,13 @@ def body_com_lin_vel_w(self) -> wp.array: This quantity is the linear velocity of the articulation bodies' center of mass frame. """ - return self._get_top_from_spatial_vector(self._body_com_lin_vel_w, self.body_com_vel_w) + self._body_com_lin_vel_w = self._get_top_from_spatial_vector(self._body_com_lin_vel_w, self.body_com_vel_w.warp) + if self._body_com_lin_vel_w_ta is None: + self._body_com_lin_vel_w_ta = ProxyArray(self._body_com_lin_vel_w) + return self._body_com_lin_vel_w_ta @property - def body_com_ang_vel_w(self) -> wp.array: + def body_com_ang_vel_w(self) -> ProxyArray: """Angular velocity of all bodies in simulation world frame. Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to @@ -1178,10 +1239,15 @@ def body_com_ang_vel_w(self) -> wp.array: This quantity is the angular velocity of the articulation bodies' center of mass frame. """ - return self._get_bottom_from_spatial_vector(self._body_com_ang_vel_w, self.body_com_vel_w) + self._body_com_ang_vel_w = self._get_bottom_from_spatial_vector( + self._body_com_ang_vel_w, self.body_com_vel_w.warp + ) + if self._body_com_ang_vel_w_ta is None: + self._body_com_ang_vel_w_ta = ProxyArray(self._body_com_ang_vel_w) + return self._body_com_ang_vel_w_ta @property - def body_com_lin_acc_w(self) -> wp.array: + def body_com_lin_acc_w(self) -> ProxyArray: """Linear acceleration of all bodies in simulation world frame. Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to @@ -1189,10 +1255,13 @@ def body_com_lin_acc_w(self) -> wp.array: This quantity is the linear acceleration of the articulation bodies' center of mass frame. """ - return self._get_top_from_spatial_vector(self._body_com_lin_acc_w, self.body_com_acc_w) + self._body_com_lin_acc_w = self._get_top_from_spatial_vector(self._body_com_lin_acc_w, self.body_com_acc_w.warp) + if self._body_com_lin_acc_w_ta is None: + self._body_com_lin_acc_w_ta = ProxyArray(self._body_com_lin_acc_w) + return self._body_com_lin_acc_w_ta @property - def body_com_ang_acc_w(self) -> wp.array: + def body_com_ang_acc_w(self) -> ProxyArray: """Angular acceleration of all bodies in simulation world frame. Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to @@ -1200,10 +1269,15 @@ def body_com_ang_acc_w(self) -> wp.array: This quantity is the angular acceleration of the articulation bodies' center of mass frame. """ - return self._get_bottom_from_spatial_vector(self._body_com_ang_acc_w, self.body_com_acc_w) + self._body_com_ang_acc_w = self._get_bottom_from_spatial_vector( + self._body_com_ang_acc_w, self.body_com_acc_w.warp + ) + if self._body_com_ang_acc_w_ta is None: + self._body_com_ang_acc_w_ta = ProxyArray(self._body_com_ang_acc_w) + return self._body_com_ang_acc_w_ta @property - def body_com_quat_b(self) -> wp.array: + def body_com_quat_b(self) -> ProxyArray: """Orientation (x, y, z, w) of the principal axes of inertia of all of the bodies in their respective link frames. @@ -1212,7 +1286,10 @@ def body_com_quat_b(self) -> wp.array: This quantity is the orientation of the principal axes of inertia relative to its body's link frame. """ - return self._get_quat_from_transform(self._body_com_quat_b, self.body_com_pose_b) + self._body_com_quat_b = self._get_quat_from_transform(self._body_com_quat_b, self.body_com_pose_b.warp) + if self._body_com_quat_b_ta is None: + self._body_com_quat_b_ta = ProxyArray(self._body_com_quat_b) + return self._body_com_quat_b_ta def _create_simulation_bindings(self) -> None: """Create simulation bindings for the root data. @@ -1339,6 +1416,11 @@ def _create_simulation_bindings(self) -> None: (self._num_instances, 0), dtype=wp.float32, device=self.device ) + # Re-pin ProxyArray wrappers to the newly created sim bindings. + # On first init, _create_buffers() handles this after all buffers exist. + if hasattr(self, "_root_link_pose_w_ta"): + self._pin_proxy_arrays() + def _create_buffers(self) -> None: """Create buffers for the root data.""" super()._create_buffers() @@ -1477,6 +1559,146 @@ def _create_buffers(self) -> None: self._body_com_ang_acc_w = None self._default_root_state = None + # Pin all ProxyArray wrappers to current buffers. + self._pin_proxy_arrays() + + def _pin_proxy_arrays(self) -> None: + """Create or rebind all pinned ProxyArray wrappers. + + Called from :meth:`_create_buffers` on first initialization and from + :meth:`_create_simulation_bindings` after a full simulation reset when + the solver recreates its internal arrays. + """ + is_rebind = hasattr(self, "_root_link_pose_w_ta") + + if is_rebind: + # Rebind sim-bound ProxyArrays to new solver arrays + self._root_link_pose_w_ta = ProxyArray(self._sim_bind_root_link_pose_w) + self._root_com_vel_w_ta = ProxyArray(self._sim_bind_root_com_vel_w) + self._body_link_pose_w_ta = ProxyArray(self._sim_bind_body_link_pose_w) + self._body_com_vel_w_ta = ProxyArray(self._sim_bind_body_com_vel_w) + self._joint_pos_ta = ProxyArray(self._sim_bind_joint_pos) + self._joint_vel_ta = ProxyArray(self._sim_bind_joint_vel) + self._joint_stiffness_ta = ProxyArray(self._sim_bind_joint_stiffness_sim) + self._joint_damping_ta = ProxyArray(self._sim_bind_joint_damping_sim) + self._joint_armature_ta = ProxyArray(self._sim_bind_joint_armature) + self._joint_friction_coeff_ta = ProxyArray(self._sim_bind_joint_friction_coeff) + self._joint_pos_limits_lower_ta = ProxyArray(self._sim_bind_joint_pos_limits_lower) + self._joint_pos_limits_upper_ta = ProxyArray(self._sim_bind_joint_pos_limits_upper) + self._joint_vel_limits_ta = ProxyArray(self._sim_bind_joint_vel_limits_sim) + self._joint_effort_limits_ta = ProxyArray(self._sim_bind_joint_effort_limits_sim) + self._body_mass_ta = ProxyArray(self._sim_bind_body_mass) + self._body_inertia_ta = ProxyArray(self._sim_bind_body_inertia) + self._body_com_pos_b_ta = ProxyArray(self._sim_bind_body_com_pos_b) + else: + # First-time creation: pin ProxyArrays to current buffers + # Category 1: sim-bound and pre-allocated buffers + # Sim-bound pointers are re-created on full reset; _create_simulation_bindings() + # calls rebind() on each ProxyArray to keep them in sync. + self._root_link_pose_w_ta = ProxyArray(self._sim_bind_root_link_pose_w) + self._root_com_vel_w_ta = ProxyArray(self._sim_bind_root_com_vel_w) + self._body_link_pose_w_ta = ProxyArray(self._sim_bind_body_link_pose_w) + self._body_com_vel_w_ta = ProxyArray(self._sim_bind_body_com_vel_w) + self._joint_pos_ta = ProxyArray(self._sim_bind_joint_pos) + self._joint_vel_ta = ProxyArray(self._sim_bind_joint_vel) + self._default_root_pose_ta = ProxyArray(self._default_root_pose) + self._default_root_vel_ta = ProxyArray(self._default_root_vel) + self._default_joint_pos_ta = ProxyArray(self._default_joint_pos) + self._default_joint_vel_ta = ProxyArray(self._default_joint_vel) + self._joint_pos_target_ta = ProxyArray(self._joint_pos_target) + self._joint_vel_target_ta = ProxyArray(self._joint_vel_target) + self._joint_effort_target_ta = ProxyArray(self._joint_effort_target) + self._computed_torque_ta = ProxyArray(self._computed_torque) + self._applied_torque_ta = ProxyArray(self._applied_torque) + self._joint_stiffness_ta = ProxyArray(self._sim_bind_joint_stiffness_sim) + self._joint_damping_ta = ProxyArray(self._sim_bind_joint_damping_sim) + self._joint_armature_ta = ProxyArray(self._sim_bind_joint_armature) + self._joint_friction_coeff_ta = ProxyArray(self._sim_bind_joint_friction_coeff) + self._joint_pos_limits_lower_ta = ProxyArray(self._sim_bind_joint_pos_limits_lower) + self._joint_pos_limits_upper_ta = ProxyArray(self._sim_bind_joint_pos_limits_upper) + self._joint_vel_limits_ta = ProxyArray(self._sim_bind_joint_vel_limits_sim) + self._joint_effort_limits_ta = ProxyArray(self._sim_bind_joint_effort_limits_sim) + self._soft_joint_pos_limits_ta = ProxyArray(self._soft_joint_pos_limits) + self._soft_joint_vel_limits_ta = ProxyArray(self._soft_joint_vel_limits) + self._gear_ratio_ta = ProxyArray(self._gear_ratio) + self._body_mass_ta = ProxyArray(self._sim_bind_body_mass) + self._body_inertia_ta = ProxyArray(self._sim_bind_body_inertia) + self._body_com_pos_b_ta = ProxyArray(self._sim_bind_body_com_pos_b) + + # Category 2: TimestampedBuffer properties + self._root_link_vel_w_ta = ProxyArray(self._root_link_vel_w.data) + self._body_link_vel_w_ta = ProxyArray(self._body_link_vel_w.data) + self._root_com_pose_w_ta = ProxyArray(self._root_com_pose_w.data) + self._body_com_pose_w_ta = ProxyArray(self._body_com_pose_w.data) + self._body_com_acc_w_ta = ProxyArray(self._body_com_acc_w.data) + self._body_com_pose_b_ta = ProxyArray(self._body_com_pose_b.data) + self._projected_gravity_b_ta = ProxyArray(self._projected_gravity_b.data) + self._heading_w_ta = ProxyArray(self._heading_w.data) + self._joint_acc_ta = ProxyArray(self._joint_acc.data) + + # -- deprecated state properties (lazy); type annotations declared once here + self._root_state_w_ta: ProxyArray | None = None + self._root_link_state_w_ta: ProxyArray | None = None + self._root_com_state_w_ta: ProxyArray | None = None + self._default_root_state_ta: ProxyArray | None = None + self._body_state_w_ta: ProxyArray | None = None + self._body_link_state_w_ta: ProxyArray | None = None + self._body_com_state_w_ta: ProxyArray | None = None + + # Invalidate lazy sliced ProxyArrays AND their backing wp.arrays so they are + # re-created from fresh data on next access. On first init the backing fields + # are already None (set by _create_buffers), so the assignments below are + # harmless no-ops. On rebind they reset stale pointers into freed transform + # memory after a sim reset. + self._root_link_pos_w_ta: ProxyArray | None = None + self._root_link_pos_w = None + self._root_link_quat_w_ta: ProxyArray | None = None + self._root_link_quat_w = None + self._root_link_lin_vel_w_ta: ProxyArray | None = None + self._root_link_lin_vel_w = None + self._root_link_ang_vel_w_ta: ProxyArray | None = None + self._root_link_ang_vel_w = None + self._root_com_pos_w_ta: ProxyArray | None = None + self._root_com_pos_w = None + self._root_com_quat_w_ta: ProxyArray | None = None + self._root_com_quat_w = None + self._root_com_lin_vel_w_ta: ProxyArray | None = None + self._root_com_lin_vel_w = None + self._root_com_ang_vel_w_ta: ProxyArray | None = None + self._root_com_ang_vel_w = None + self._body_link_pos_w_ta: ProxyArray | None = None + self._body_link_pos_w = None + self._body_link_quat_w_ta: ProxyArray | None = None + self._body_link_quat_w = None + self._body_link_lin_vel_w_ta: ProxyArray | None = None + self._body_link_lin_vel_w = None + self._body_link_ang_vel_w_ta: ProxyArray | None = None + self._body_link_ang_vel_w = None + self._body_com_pos_w_ta: ProxyArray | None = None + self._body_com_pos_w = None + self._body_com_quat_w_ta: ProxyArray | None = None + self._body_com_quat_w = None + self._body_com_lin_vel_w_ta: ProxyArray | None = None + self._body_com_lin_vel_w = None + self._body_com_ang_vel_w_ta: ProxyArray | None = None + self._body_com_ang_vel_w = None + self._body_com_lin_acc_w_ta: ProxyArray | None = None + self._body_com_lin_acc_w = None + self._body_com_ang_acc_w_ta: ProxyArray | None = None + self._body_com_ang_acc_w = None + self._body_com_quat_b_ta: ProxyArray | None = None + self._body_com_quat_b = None + self._joint_pos_limits_ta: ProxyArray | None = None + self._joint_pos_limits = None + self._root_link_lin_vel_b_ta: ProxyArray | None = None + self._root_link_lin_vel_b = None + self._root_link_ang_vel_b_ta: ProxyArray | None = None + self._root_link_ang_vel_b = None + self._root_com_lin_vel_b_ta: ProxyArray | None = None + self._root_com_lin_vel_b = None + self._root_com_ang_vel_b_ta: ProxyArray | None = None + self._root_com_ang_vel_b = None + """ Internal helpers. """ @@ -1681,7 +1903,7 @@ def _get_bottom_from_spatial_vector(self, source: wp.array | None, spatial_vecto """ @property - def root_state_w(self) -> wp.array: + def root_state_w(self) -> ProxyArray: """Deprecated, same as :attr:`root_link_pose_w` and :attr:`root_com_vel_w`.""" warnings.warn( "The `root_state_w` property will be deprecated in a IsaacLab 4.0. Please use `root_link_pose_w` and " @@ -1693,13 +1915,14 @@ def root_state_w(self) -> wp.array: self._root_state_w = TimestampedBuffer( shape=(self._num_instances,), dtype=shared_kernels.vec13f, device=self.device ) + self._root_state_w_ta = ProxyArray(self._root_state_w.data) if self._root_state_w.timestamp < self._sim_timestamp: wp.launch( shared_kernels.concat_root_pose_and_vel_to_state, dim=(self._num_instances), inputs=[ - self.root_link_pose_w, - self.root_com_vel_w, + self.root_link_pose_w.warp, + self.root_com_vel_w.warp, ], outputs=[ self._root_state_w.data, @@ -1708,10 +1931,10 @@ def root_state_w(self) -> wp.array: ) self._root_state_w.timestamp = self._sim_timestamp - return self._root_state_w.data + return self._root_state_w_ta @property - def root_link_state_w(self) -> wp.array: + def root_link_state_w(self) -> ProxyArray: """Deprecated, same as :attr:`root_link_pose_w` and :attr:`root_link_vel_w`.""" warnings.warn( "The `root_link_state_w` property will be deprecated in a IsaacLab 4.0. Please use `root_link_pose_w` and " @@ -1723,13 +1946,14 @@ def root_link_state_w(self) -> wp.array: self._root_link_state_w = TimestampedBuffer( shape=(self._num_instances,), dtype=shared_kernels.vec13f, device=self.device ) + self._root_link_state_w_ta = ProxyArray(self._root_link_state_w.data) if self._root_link_state_w.timestamp < self._sim_timestamp: wp.launch( shared_kernels.concat_root_pose_and_vel_to_state, dim=self._num_instances, inputs=[ - self.root_link_pose_w, - self.root_link_vel_w, + self.root_link_pose_w.warp, + self.root_link_vel_w.warp, ], outputs=[ self._root_link_state_w.data, @@ -1738,10 +1962,10 @@ def root_link_state_w(self) -> wp.array: ) self._root_link_state_w.timestamp = self._sim_timestamp - return self._root_link_state_w.data + return self._root_link_state_w_ta @property - def root_com_state_w(self) -> wp.array: + def root_com_state_w(self) -> ProxyArray: """Deprecated, same as :attr:`root_com_pose_w` and :attr:`root_com_vel_w`.""" warnings.warn( "The `root_com_state_w` property will be deprecated in a IsaacLab 4.0. Please use `root_com_pose_w` and " @@ -1753,13 +1977,14 @@ def root_com_state_w(self) -> wp.array: self._root_com_state_w = TimestampedBuffer( shape=(self._num_instances,), dtype=shared_kernels.vec13f, device=self.device ) + self._root_com_state_w_ta = ProxyArray(self._root_com_state_w.data) if self._root_com_state_w.timestamp < self._sim_timestamp: wp.launch( shared_kernels.concat_root_pose_and_vel_to_state, dim=self._num_instances, inputs=[ - self.root_com_pose_w, - self.root_com_vel_w, + self.root_com_pose_w.warp, + self.root_com_vel_w.warp, ], outputs=[ self._root_com_state_w.data, @@ -1768,10 +1993,10 @@ def root_com_state_w(self) -> wp.array: ) self._root_com_state_w.timestamp = self._sim_timestamp - return self._root_com_state_w.data + return self._root_com_state_w_ta @property - def default_root_state(self) -> wp.array: + def default_root_state(self) -> ProxyArray: """Default root state ``[pos, quat, lin_vel, ang_vel]`` in the local environment frame. @@ -1788,6 +2013,7 @@ def default_root_state(self) -> wp.array: ) if self._default_root_state is None: self._default_root_state = wp.zeros((self._num_instances), dtype=shared_kernels.vec13f, device=self.device) + self._default_root_state_ta = ProxyArray(self._default_root_state) wp.launch( shared_kernels.concat_root_pose_and_vel_to_state, dim=self._num_instances, @@ -1800,10 +2026,10 @@ def default_root_state(self) -> wp.array: ], device=self.device, ) - return self._default_root_state + return self._default_root_state_ta @property - def body_state_w(self): + def body_state_w(self) -> ProxyArray: """State of all bodies `[pos, quat, lin_vel, ang_vel]` in simulation world frame. Shape is (num_instances, num_bodies, 13). @@ -1820,13 +2046,14 @@ def body_state_w(self): self._body_state_w = TimestampedBuffer( (self._num_instances, self._num_bodies), self.device, shared_kernels.vec13f ) + self._body_state_w_ta = ProxyArray(self._body_state_w.data) if self._body_state_w.timestamp < self._sim_timestamp: wp.launch( shared_kernels.concat_body_pose_and_vel_to_state, dim=(self._num_instances, self._num_bodies), inputs=[ - self.body_link_pose_w, - self.body_com_vel_w, + self.body_link_pose_w.warp, + self.body_com_vel_w.warp, ], outputs=[ self._body_state_w.data, @@ -1835,10 +2062,10 @@ def body_state_w(self): ) self._body_state_w.timestamp = self._sim_timestamp - return self._body_state_w.data + return self._body_state_w_ta @property - def body_link_state_w(self): + def body_link_state_w(self) -> ProxyArray: """State of all bodies' link frame`[pos, quat, lin_vel, ang_vel]` in simulation world frame. Shape is (num_instances, num_bodies, 13). @@ -1854,13 +2081,14 @@ def body_link_state_w(self): self._body_link_state_w = TimestampedBuffer( (self._num_instances, self._num_bodies), self.device, shared_kernels.vec13f ) + self._body_link_state_w_ta = ProxyArray(self._body_link_state_w.data) if self._body_link_state_w.timestamp < self._sim_timestamp: wp.launch( shared_kernels.concat_body_pose_and_vel_to_state, dim=(self._num_instances, self._num_bodies), inputs=[ - self.body_link_pose_w, - self.body_link_vel_w, + self.body_link_pose_w.warp, + self.body_link_vel_w.warp, ], outputs=[ self._body_link_state_w.data, @@ -1869,10 +2097,10 @@ def body_link_state_w(self): ) self._body_link_state_w.timestamp = self._sim_timestamp - return self._body_link_state_w.data + return self._body_link_state_w_ta @property - def body_com_state_w(self): + def body_com_state_w(self) -> ProxyArray: """State of all bodies center of mass `[pos, quat, lin_vel, ang_vel]` in simulation world frame. Shape is (num_instances, num_bodies, 13). @@ -1890,13 +2118,14 @@ def body_com_state_w(self): self._body_com_state_w = TimestampedBuffer( (self._num_instances, self._num_bodies), self.device, shared_kernels.vec13f ) + self._body_com_state_w_ta = ProxyArray(self._body_com_state_w.data) if self._body_com_state_w.timestamp < self._sim_timestamp: wp.launch( shared_kernels.concat_body_pose_and_vel_to_state, dim=(self._num_instances, self._num_bodies), inputs=[ - self.body_com_pose_w, - self.body_com_vel_w, + self.body_com_pose_w.warp, + self.body_com_vel_w.warp, ], outputs=[ self._body_com_state_w.data, @@ -1905,4 +2134,4 @@ def body_com_state_w(self): ) self._body_com_state_w.timestamp = self._sim_timestamp - return self._body_com_state_w.data + return self._body_com_state_w_ta diff --git a/source/isaaclab_newton/isaaclab_newton/assets/rigid_object/rigid_object_data.py b/source/isaaclab_newton/isaaclab_newton/assets/rigid_object/rigid_object_data.py index 514d4e836d34..0e9ecc8a41d0 100644 --- a/source/isaaclab_newton/isaaclab_newton/assets/rigid_object/rigid_object_data.py +++ b/source/isaaclab_newton/isaaclab_newton/assets/rigid_object/rigid_object_data.py @@ -15,6 +15,7 @@ from isaaclab.assets.rigid_object.base_rigid_object_data import BaseRigidObjectData from isaaclab.utils.buffers import TimestampedBufferWarp as TimestampedBuffer from isaaclab.utils.math import normalize +from isaaclab.utils.warp import ProxyArray from isaaclab.utils.warp.utils import capture_unsafe from isaaclab_newton.assets import kernels as shared_kernels @@ -84,8 +85,8 @@ def __init__(self, root_view: ArticulationView, device: str): forward_vec = torch.tensor((1.0, 0.0, 0.0), device=self.device).repeat(self._root_view.count, 1) # Initialize constants - self.GRAVITY_VEC_W = wp.from_torch(gravity_dir, dtype=wp.vec3f) - self.FORWARD_VEC_B = wp.from_torch(forward_vec, dtype=wp.vec3f) + self.GRAVITY_VEC_W = ProxyArray(wp.from_torch(gravity_dir, dtype=wp.vec3f)) + self.FORWARD_VEC_B = ProxyArray(wp.from_torch(forward_vec, dtype=wp.vec3f)) self._create_simulation_bindings() self._create_buffers() @@ -136,13 +137,13 @@ def update(self, dt: float) -> None: """ @property - def default_root_pose(self) -> wp.array: + def default_root_pose(self) -> ProxyArray: """Default root pose ``[pos, quat]`` in local environment frame. Shape is (num_instances,), dtype = wp.transformf. In torch this resolves to (num_instances, 7). The position and quaternion are of the rigid body's actor frame. """ - return self._default_root_pose + return self._default_root_pose_ta @default_root_pose.setter def default_root_pose(self, value: wp.array) -> None: @@ -159,13 +160,13 @@ def default_root_pose(self, value: wp.array) -> None: self._default_root_pose.assign(value) @property - def default_root_vel(self) -> wp.array: + def default_root_vel(self) -> ProxyArray: """Default root velocity ``[lin_vel, ang_vel]`` in local environment frame. Shape is (num_instances,), dtype = wp.spatial_vectorf. In torch this resolves to (num_instances, 6). The linear and angular velocities are of the rigid body's center of mass frame. """ - return self._default_root_vel + return self._default_root_vel_ta @default_root_vel.setter def default_root_vel(self, value: wp.array) -> None: @@ -186,18 +187,18 @@ def default_root_vel(self, value: wp.array) -> None: """ @property - def root_link_pose_w(self) -> wp.array: + def root_link_pose_w(self) -> ProxyArray: """Root link pose ``[pos, quat]`` in simulation world frame. Shape is (num_instances,), dtype = wp.transformf. In torch this resolves to (num_instances, 7). This quantity is the pose of the actor frame of the root rigid body relative to the world. The orientation is provided in (x, y, z, w) format. """ - return self._sim_bind_root_link_pose_w + return self._root_link_pose_w_ta @property @capture_unsafe(_LAZY_CAPTURE_REASON) - def root_link_vel_w(self) -> wp.array: + def root_link_vel_w(self) -> ProxyArray: """Root link velocity ``[lin_vel, ang_vel]`` in simulation world frame. Shape is (num_instances,), dtype = wp.spatial_vectorf. In torch this resolves to (num_instances, 6). @@ -210,9 +211,9 @@ def root_link_vel_w(self) -> wp.array: shared_kernels.get_root_link_vel_from_root_com_vel, dim=self._num_instances, inputs=[ - self.root_com_vel_w, - self.root_link_pose_w, - self.body_com_pos_b, + self.root_com_vel_w.warp, + self.root_link_pose_w.warp, + self.body_com_pos_b.warp, ], outputs=[ self._root_link_vel_w.data, @@ -221,11 +222,11 @@ def root_link_vel_w(self) -> wp.array: ) self._root_link_vel_w.timestamp = self._sim_timestamp - return self._root_link_vel_w.data + return self._root_link_vel_w_ta @property @capture_unsafe(_LAZY_CAPTURE_REASON) - def root_com_pose_w(self) -> wp.array: + def root_com_pose_w(self) -> ProxyArray: """Root center of mass pose ``[pos, quat]`` in simulation world frame. Shape is (num_instances,), dtype = wp.transformf. In torch this resolves to (num_instances, 7). @@ -238,8 +239,8 @@ def root_com_pose_w(self) -> wp.array: shared_kernels.get_root_com_pose_from_root_link_pose, dim=self._num_instances, inputs=[ - self.root_link_pose_w, - self.body_com_pos_b, + self.root_link_pose_w.warp, + self.body_com_pos_b.warp, ], outputs=[ self._root_com_pose_w.data, @@ -248,84 +249,88 @@ def root_com_pose_w(self) -> wp.array: ) self._root_com_pose_w.timestamp = self._sim_timestamp - return self._root_com_pose_w.data + return self._root_com_pose_w_ta @property - def root_com_vel_w(self) -> wp.array: + def root_com_vel_w(self) -> ProxyArray: """Root center of mass velocity ``[lin_vel, ang_vel]`` in simulation world frame. Shape is (num_instances,), dtype = wp.spatial_vectorf. In torch this resolves to (num_instances, 6). This quantity contains the linear and angular velocities of the root rigid body's center of mass frame relative to the world. """ - return self._sim_bind_root_com_vel_w + return self._root_com_vel_w_ta """ Body state properties. """ @property - def body_mass(self) -> wp.array: + def body_mass(self) -> ProxyArray: """Mass of all bodies in the simulation world frame. Shape is (num_instances, 1, 1), dtype = wp.float32. In torch this resolves to (num_instances, 1, 1). """ - return self._sim_bind_body_mass + return self._body_mass_ta @property - def body_inertia(self) -> wp.array: + def body_inertia(self) -> ProxyArray: """Inertia of all bodies in the simulation world frame. Shape is (num_instances, 1, 9), dtype = wp.float32. In torch this resolves to (num_instances, 1, 9). """ - return self._sim_bind_body_inertia + return self._body_inertia_ta @property - def body_link_pose_w(self) -> wp.array: + def body_link_pose_w(self) -> ProxyArray: """Body link pose ``[pos, quat]`` in simulation world frame. Shape is (num_instances, 1), dtype = wp.transformf. In torch this resolves to (num_instances, 1, 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. """ - return self._sim_bind_body_link_pose_w + return self._body_link_pose_w_ta @property @capture_unsafe(_LAZY_CAPTURE_REASON) - def body_link_vel_w(self) -> wp.array: + def body_link_vel_w(self) -> ProxyArray: """Body link velocity ``[lin_vel, ang_vel]`` in simulation world frame. Shape is (num_instances, 1), dtype = wp.spatial_vectorf. In torch this resolves to (num_instances, 1, 6). This quantity contains the linear and angular velocities of the actor frame of the root rigid body relative to the world. """ - return self.root_link_vel_w.reshape((self._num_instances, 1)) + if self._body_link_vel_w_ta is None: + self._body_link_vel_w_ta = ProxyArray(self.root_link_vel_w.warp.reshape((self._num_instances, 1))) + return self._body_link_vel_w_ta @property @capture_unsafe(_LAZY_CAPTURE_REASON) - def body_com_pose_w(self) -> wp.array: + def body_com_pose_w(self) -> ProxyArray: """Body center of mass pose ``[pos, quat]`` in simulation world frame. Shape is (num_instances, 1), dtype = wp.transformf. In torch this resolves to (num_instances, 1, 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. """ - return self.root_com_pose_w.reshape((self._num_instances, 1)) + if self._body_com_pose_w_ta is None: + self._body_com_pose_w_ta = ProxyArray(self.root_com_pose_w.warp.reshape((self._num_instances, 1))) + return self._body_com_pose_w_ta @property - def body_com_vel_w(self) -> wp.array: + def body_com_vel_w(self) -> ProxyArray: """Body center of mass velocity ``[lin_vel, ang_vel]`` in simulation world frame. Shape is (num_instances, 1), dtype = wp.spatial_vectorf. In torch this resolves to (num_instances, 1, 6). This quantity contains the linear and angular velocities of the root rigid body's center of mass frame relative to the world. """ - return self._sim_bind_body_com_vel_w + return self._body_com_vel_w_ta @property - def body_com_acc_w(self) -> wp.array: + def body_com_acc_w(self) -> ProxyArray: """Acceleration of all bodies ``[lin_acc, ang_acc]`` in the simulation world frame. Shape is (num_instances, 1), dtype = wp.spatial_vectorf. In torch this resolves to (num_instances, 1, 6). @@ -348,19 +353,19 @@ def body_com_acc_w(self) -> wp.array: # set the buffer data and timestamp self._body_com_acc_w.timestamp = self._sim_timestamp # update the previous velocity - return self._body_com_acc_w.data + return self._body_com_acc_w_ta @property - def body_com_pos_b(self) -> wp.array: + def body_com_pos_b(self) -> ProxyArray: """Center of mass position of all of the bodies in their respective link frames. Shape is (num_instances, 1), dtype = wp.vec3f. In torch this resolves to (num_instances, 1, 3). This quantity is the center of mass location relative to its body's link frame. """ - return self._sim_bind_body_com_pos_b + return self._body_com_pos_b_ta @property - def body_com_pose_b(self) -> wp.array: + def body_com_pose_b(self) -> ProxyArray: """Center of mass pose ``[pos, quat]`` of all bodies in their respective body's link frames. Shape is (num_instances, 1), dtype = wp.transformf. In torch this resolves to (num_instances, 1, 7). @@ -379,7 +384,7 @@ def body_com_pose_b(self) -> wp.array: shared_kernels.make_dummy_body_com_pose_b, dim=(self._num_instances, 1), inputs=[ - self.body_com_pos_b, + self.body_com_pos_b.warp, ], outputs=[ self._body_com_pose_b.data, @@ -387,7 +392,7 @@ def body_com_pose_b(self) -> wp.array: device=self.device, ) self._body_com_pose_b.timestamp = self._sim_timestamp - return self._body_com_pose_b.data + return self._body_com_pose_b_ta """ Derived Properties. @@ -395,7 +400,7 @@ def body_com_pose_b(self) -> wp.array: @property @capture_unsafe(_LAZY_CAPTURE_REASON) - def projected_gravity_b(self) -> wp.array: + def projected_gravity_b(self) -> ProxyArray: """Projection of the gravity direction on base frame. Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). @@ -404,16 +409,16 @@ def projected_gravity_b(self) -> wp.array: wp.launch( shared_kernels.quat_apply_inverse_1D_kernel, dim=self._num_instances, - inputs=[self.GRAVITY_VEC_W, self.root_link_quat_w], + inputs=[self.GRAVITY_VEC_W.warp, self.root_link_quat_w.warp], outputs=[self._projected_gravity_b.data], device=self.device, ) self._projected_gravity_b.timestamp = self._sim_timestamp - return self._projected_gravity_b.data + return self._projected_gravity_b_ta @property @capture_unsafe(_LAZY_CAPTURE_REASON) - def heading_w(self) -> wp.array: + def heading_w(self) -> ProxyArray: """Yaw heading of the base frame (in radians). Shape is (num_instances,), dtype = wp.float32. In torch this resolves to (num_instances,). @@ -426,16 +431,16 @@ def heading_w(self) -> wp.array: wp.launch( shared_kernels.root_heading_w, dim=self._num_instances, - inputs=[self.FORWARD_VEC_B, self.root_link_quat_w], + inputs=[self.FORWARD_VEC_B.warp, self.root_link_quat_w.warp], outputs=[self._heading_w.data], device=self.device, ) self._heading_w.timestamp = self._sim_timestamp - return self._heading_w.data + return self._heading_w_ta @property @capture_unsafe(_LAZY_CAPTURE_REASON) - def root_link_lin_vel_b(self) -> wp.array: + def root_link_lin_vel_b(self) -> ProxyArray: """Root link linear velocity in base frame. Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). @@ -446,19 +451,20 @@ def root_link_lin_vel_b(self) -> wp.array: self._root_link_lin_vel_b = TimestampedBuffer( shape=(self._num_instances,), dtype=wp.vec3f, device=self.device ) + self._root_link_lin_vel_b_ta = ProxyArray(self._root_link_lin_vel_b.data) if self._root_link_lin_vel_b.timestamp < self._sim_timestamp: wp.launch( shared_kernels.quat_apply_inverse_1D_kernel, dim=self._num_instances, - inputs=[self.root_link_lin_vel_w, self.root_link_quat_w], + inputs=[self.root_link_lin_vel_w.warp, self.root_link_quat_w.warp], outputs=[self._root_link_lin_vel_b.data], device=self.device, ) self._root_link_lin_vel_b.timestamp = self._sim_timestamp - return self._root_link_lin_vel_b.data + return self._root_link_lin_vel_b_ta @property - def root_link_ang_vel_b(self) -> wp.array: + def root_link_ang_vel_b(self) -> ProxyArray: """Root link angular velocity in base frame. Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). @@ -469,19 +475,20 @@ def root_link_ang_vel_b(self) -> wp.array: self._root_link_ang_vel_b = TimestampedBuffer( shape=(self._num_instances,), dtype=wp.vec3f, device=self.device ) + self._root_link_ang_vel_b_ta = ProxyArray(self._root_link_ang_vel_b.data) if self._root_link_ang_vel_b.timestamp < self._sim_timestamp: wp.launch( shared_kernels.quat_apply_inverse_1D_kernel, dim=self._num_instances, - inputs=[self.root_link_ang_vel_w, self.root_link_quat_w], + inputs=[self.root_link_ang_vel_w.warp, self.root_link_quat_w.warp], outputs=[self._root_link_ang_vel_b.data], device=self.device, ) self._root_link_ang_vel_b.timestamp = self._sim_timestamp - return self._root_link_ang_vel_b.data + return self._root_link_ang_vel_b_ta @property - def root_com_lin_vel_b(self) -> wp.array: + def root_com_lin_vel_b(self) -> ProxyArray: """Root center of mass linear velocity in base frame. Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). @@ -492,19 +499,20 @@ def root_com_lin_vel_b(self) -> wp.array: self._root_com_lin_vel_b = TimestampedBuffer( shape=(self._num_instances,), dtype=wp.vec3f, device=self.device ) + self._root_com_lin_vel_b_ta = ProxyArray(self._root_com_lin_vel_b.data) if self._root_com_lin_vel_b.timestamp < self._sim_timestamp: wp.launch( shared_kernels.quat_apply_inverse_1D_kernel, dim=self._num_instances, - inputs=[self.root_com_lin_vel_w, self.root_link_quat_w], + inputs=[self.root_com_lin_vel_w.warp, self.root_link_quat_w.warp], outputs=[self._root_com_lin_vel_b.data], device=self.device, ) self._root_com_lin_vel_b.timestamp = self._sim_timestamp - return self._root_com_lin_vel_b.data + return self._root_com_lin_vel_b_ta @property - def root_com_ang_vel_b(self) -> wp.array: + def root_com_ang_vel_b(self) -> ProxyArray: """Root center of mass angular velocity in base frame. Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). @@ -515,192 +523,264 @@ def root_com_ang_vel_b(self) -> wp.array: self._root_com_ang_vel_b = TimestampedBuffer( shape=(self._num_instances,), dtype=wp.vec3f, device=self.device ) + self._root_com_ang_vel_b_ta = ProxyArray(self._root_com_ang_vel_b.data) if self._root_com_ang_vel_b.timestamp < self._sim_timestamp: wp.launch( shared_kernels.quat_apply_inverse_1D_kernel, dim=self._num_instances, - inputs=[self.root_com_ang_vel_w, self.root_link_quat_w], + inputs=[self.root_com_ang_vel_w.warp, self.root_link_quat_w.warp], outputs=[self._root_com_ang_vel_b.data], device=self.device, ) self._root_com_ang_vel_b.timestamp = self._sim_timestamp - return self._root_com_ang_vel_b.data + return self._root_com_ang_vel_b_ta """ Sliced properties. """ @property - def root_link_pos_w(self) -> wp.array: + def root_link_pos_w(self) -> ProxyArray: """Root link position in simulation world frame. Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). This quantity is the position of the actor frame of the root rigid body relative to the world. """ - return self._get_pos_from_transform(self._root_link_pos_w, self.root_link_pose_w) + self._root_link_pos_w = self._get_pos_from_transform(self._root_link_pos_w, self.root_link_pose_w.warp) + if self._root_link_pos_w_ta is None: + self._root_link_pos_w_ta = ProxyArray(self._root_link_pos_w) + return self._root_link_pos_w_ta @property - def root_link_quat_w(self) -> wp.array: + def root_link_quat_w(self) -> ProxyArray: """Root link orientation (x, y, z, w) in simulation world frame. Shape is (num_instances,), dtype = wp.quatf. In torch this resolves to (num_instances, 4). This quantity is the orientation of the actor frame of the root rigid body. """ - return self._get_quat_from_transform(self._root_link_quat_w, self.root_link_pose_w) + self._root_link_quat_w = self._get_quat_from_transform(self._root_link_quat_w, self.root_link_pose_w.warp) + if self._root_link_quat_w_ta is None: + self._root_link_quat_w_ta = ProxyArray(self._root_link_quat_w) + return self._root_link_quat_w_ta @property - def root_link_lin_vel_w(self) -> wp.array: + def root_link_lin_vel_w(self) -> ProxyArray: """Root linear velocity in simulation world frame. Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). This quantity is the linear velocity of the root rigid body's actor frame relative to the world. """ - return self._get_top_from_spatial_vector(self._root_link_lin_vel_w, self.root_link_vel_w) + self._root_link_lin_vel_w = self._get_top_from_spatial_vector( + self._root_link_lin_vel_w, self.root_link_vel_w.warp + ) + if self._root_link_lin_vel_w_ta is None: + self._root_link_lin_vel_w_ta = ProxyArray(self._root_link_lin_vel_w) + return self._root_link_lin_vel_w_ta @property - def root_link_ang_vel_w(self) -> wp.array: + def root_link_ang_vel_w(self) -> ProxyArray: """Root link angular velocity in simulation world frame. Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). This quantity is the angular velocity of the actor frame of the root rigid body relative to the world. """ - return self._get_bottom_from_spatial_vector(self._root_link_ang_vel_w, self.root_link_vel_w) + self._root_link_ang_vel_w = self._get_bottom_from_spatial_vector( + self._root_link_ang_vel_w, self.root_link_vel_w.warp + ) + if self._root_link_ang_vel_w_ta is None: + self._root_link_ang_vel_w_ta = ProxyArray(self._root_link_ang_vel_w) + return self._root_link_ang_vel_w_ta @property - def root_com_pos_w(self) -> wp.array: + def root_com_pos_w(self) -> ProxyArray: """Root center of mass position in simulation world frame. Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). This quantity is the position of the center of mass frame of the root rigid body relative to the world. """ - return self._get_pos_from_transform(self._root_com_pos_w, self.root_com_pose_w) + self._root_com_pos_w = self._get_pos_from_transform(self._root_com_pos_w, self.root_com_pose_w.warp) + if self._root_com_pos_w_ta is None: + self._root_com_pos_w_ta = ProxyArray(self._root_com_pos_w) + return self._root_com_pos_w_ta @property - def root_com_quat_w(self) -> wp.array: + def root_com_quat_w(self) -> ProxyArray: """Root center of mass orientation (x, y, z, w) in simulation world frame. Shape is (num_instances,), dtype = wp.quatf. In torch this resolves to (num_instances, 4). This quantity is the orientation of the principal axes of inertia of the root rigid body relative to the world. """ - return self._get_quat_from_transform(self._root_com_quat_w, self.root_com_pose_w) + self._root_com_quat_w = self._get_quat_from_transform(self._root_com_quat_w, self.root_com_pose_w.warp) + if self._root_com_quat_w_ta is None: + self._root_com_quat_w_ta = ProxyArray(self._root_com_quat_w) + return self._root_com_quat_w_ta @property - def root_com_lin_vel_w(self) -> wp.array: + def root_com_lin_vel_w(self) -> ProxyArray: """Root center of mass linear velocity in simulation world frame. Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). This quantity is the linear velocity of the root rigid body's center of mass frame relative to the world. """ - return self._get_top_from_spatial_vector(self._root_com_lin_vel_w, self.root_com_vel_w) + self._root_com_lin_vel_w = self._get_top_from_spatial_vector(self._root_com_lin_vel_w, self.root_com_vel_w.warp) + if self._root_com_lin_vel_w_ta is None: + self._root_com_lin_vel_w_ta = ProxyArray(self._root_com_lin_vel_w) + return self._root_com_lin_vel_w_ta @property - def root_com_ang_vel_w(self) -> wp.array: + def root_com_ang_vel_w(self) -> ProxyArray: """Root center of mass angular velocity in simulation world frame. Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). This quantity is the angular velocity of the root rigid body's center of mass frame relative to the world. """ - return self._get_bottom_from_spatial_vector(self._root_com_ang_vel_w, self.root_com_vel_w) + self._root_com_ang_vel_w = self._get_bottom_from_spatial_vector( + self._root_com_ang_vel_w, self.root_com_vel_w.warp + ) + if self._root_com_ang_vel_w_ta is None: + self._root_com_ang_vel_w_ta = ProxyArray(self._root_com_ang_vel_w) + return self._root_com_ang_vel_w_ta @property - def body_link_pos_w(self) -> wp.array: + def body_link_pos_w(self) -> ProxyArray: """Positions of all bodies in simulation world frame. Shape is (num_instances, 1), dtype = wp.vec3f. In torch this resolves to (num_instances, 1, 3). This quantity is the position of the rigid bodies' actor frame relative to the world. """ - return self._get_pos_from_transform(self._body_link_pos_w, self.body_link_pose_w) + self._body_link_pos_w = self._get_pos_from_transform(self._body_link_pos_w, self.body_link_pose_w.warp) + if self._body_link_pos_w_ta is None: + self._body_link_pos_w_ta = ProxyArray(self._body_link_pos_w) + return self._body_link_pos_w_ta @property - def body_link_quat_w(self) -> wp.array: + def body_link_quat_w(self) -> ProxyArray: """Orientation (x, y, z, w) of all bodies in simulation world frame. Shape is (num_instances, 1), dtype = wp.quatf. In torch this resolves to (num_instances, 1, 4). This quantity is the orientation of the rigid bodies' actor frame relative to the world. """ - return self._get_quat_from_transform(self._body_link_quat_w, self.body_link_pose_w) + self._body_link_quat_w = self._get_quat_from_transform(self._body_link_quat_w, self.body_link_pose_w.warp) + if self._body_link_quat_w_ta is None: + self._body_link_quat_w_ta = ProxyArray(self._body_link_quat_w) + return self._body_link_quat_w_ta @property - def body_link_lin_vel_w(self) -> wp.array: + def body_link_lin_vel_w(self) -> ProxyArray: """Linear velocity of all bodies in simulation world frame. Shape is (num_instances, 1), dtype = wp.vec3f. In torch this resolves to (num_instances, 1, 3). This quantity is the linear velocity of the rigid bodies' actor frame relative to the world. """ - return self._get_top_from_spatial_vector(self._body_link_lin_vel_w, self.body_link_vel_w) + self._body_link_lin_vel_w = self._get_top_from_spatial_vector( + self._body_link_lin_vel_w, self.body_link_vel_w.warp + ) + if self._body_link_lin_vel_w_ta is None: + self._body_link_lin_vel_w_ta = ProxyArray(self._body_link_lin_vel_w) + return self._body_link_lin_vel_w_ta @property - def body_link_ang_vel_w(self) -> wp.array: + def body_link_ang_vel_w(self) -> ProxyArray: """Angular velocity of all bodies in simulation world frame. Shape is (num_instances, 1), dtype = wp.vec3f. In torch this resolves to (num_instances, 1, 3). This quantity is the angular velocity of the rigid bodies' actor frame relative to the world. """ - return self._get_bottom_from_spatial_vector(self._body_link_ang_vel_w, self.body_link_vel_w) + self._body_link_ang_vel_w = self._get_bottom_from_spatial_vector( + self._body_link_ang_vel_w, self.body_link_vel_w.warp + ) + if self._body_link_ang_vel_w_ta is None: + self._body_link_ang_vel_w_ta = ProxyArray(self._body_link_ang_vel_w) + return self._body_link_ang_vel_w_ta @property - def body_com_pos_w(self) -> wp.array: + def body_com_pos_w(self) -> ProxyArray: """Positions of all bodies' center of mass in simulation world frame. Shape is (num_instances, 1), dtype = wp.vec3f. In torch this resolves to (num_instances, 1, 3). This quantity is the position of the rigid bodies' center of mass frame. """ - return self._get_pos_from_transform(self._body_com_pos_w, self.body_com_pose_w) + self._body_com_pos_w = self._get_pos_from_transform(self._body_com_pos_w, self.body_com_pose_w.warp) + if self._body_com_pos_w_ta is None: + self._body_com_pos_w_ta = ProxyArray(self._body_com_pos_w) + return self._body_com_pos_w_ta @property - def body_com_quat_w(self) -> wp.array: + def body_com_quat_w(self) -> ProxyArray: """Orientation (x, y, z, w) of the principal axes of inertia of all bodies in simulation world frame. Shape is (num_instances, 1), dtype = wp.quatf. In torch this resolves to (num_instances, 1, 4). This quantity is the orientation of the principal axes of inertia of the rigid bodies. """ - return self._get_quat_from_transform(self._body_com_quat_w, self.body_com_pose_w) + self._body_com_quat_w = self._get_quat_from_transform(self._body_com_quat_w, self.body_com_pose_w.warp) + if self._body_com_quat_w_ta is None: + self._body_com_quat_w_ta = ProxyArray(self._body_com_quat_w) + return self._body_com_quat_w_ta @property - def body_com_lin_vel_w(self) -> wp.array: + def body_com_lin_vel_w(self) -> ProxyArray: """Linear velocity of all bodies in simulation world frame. Shape is (num_instances, 1), dtype = wp.vec3f. In torch this resolves to (num_instances, 1, 3). This quantity is the linear velocity of the rigid bodies' center of mass frame. """ - return self._get_top_from_spatial_vector(self._body_com_lin_vel_w, self.body_com_vel_w) + self._body_com_lin_vel_w = self._get_top_from_spatial_vector(self._body_com_lin_vel_w, self.body_com_vel_w.warp) + if self._body_com_lin_vel_w_ta is None: + self._body_com_lin_vel_w_ta = ProxyArray(self._body_com_lin_vel_w) + return self._body_com_lin_vel_w_ta @property - def body_com_ang_vel_w(self) -> wp.array: + def body_com_ang_vel_w(self) -> ProxyArray: """Angular velocity of all bodies in simulation world frame. Shape is (num_instances, 1), dtype = wp.vec3f. In torch this resolves to (num_instances, 1, 3). This quantity is the angular velocity of the rigid bodies' center of mass frame. """ - return self._get_bottom_from_spatial_vector(self._body_com_ang_vel_w, self.body_com_vel_w) + self._body_com_ang_vel_w = self._get_bottom_from_spatial_vector( + self._body_com_ang_vel_w, self.body_com_vel_w.warp + ) + if self._body_com_ang_vel_w_ta is None: + self._body_com_ang_vel_w_ta = ProxyArray(self._body_com_ang_vel_w) + return self._body_com_ang_vel_w_ta @property - def body_com_lin_acc_w(self) -> wp.array: + def body_com_lin_acc_w(self) -> ProxyArray: """Linear acceleration of all bodies in simulation world frame. Shape is (num_instances, 1), dtype = wp.vec3f. In torch this resolves to (num_instances, 1, 3). This quantity is the linear acceleration of the rigid bodies' center of mass frame. """ - return self._get_top_from_spatial_vector(self._body_com_lin_acc_w, self.body_com_acc_w) + self._body_com_lin_acc_w = self._get_top_from_spatial_vector(self._body_com_lin_acc_w, self.body_com_acc_w.warp) + if self._body_com_lin_acc_w_ta is None: + self._body_com_lin_acc_w_ta = ProxyArray(self._body_com_lin_acc_w) + return self._body_com_lin_acc_w_ta @property - def body_com_ang_acc_w(self) -> wp.array: + def body_com_ang_acc_w(self) -> ProxyArray: """Angular acceleration of all bodies in simulation world frame. Shape is (num_instances, 1), dtype = wp.vec3f. In torch this resolves to (num_instances, 1, 3). This quantity is the angular acceleration of the rigid bodies' center of mass frame. """ - return self._get_bottom_from_spatial_vector(self._body_com_ang_acc_w, self.body_com_acc_w) + self._body_com_ang_acc_w = self._get_bottom_from_spatial_vector( + self._body_com_ang_acc_w, self.body_com_acc_w.warp + ) + if self._body_com_ang_acc_w_ta is None: + self._body_com_ang_acc_w_ta = ProxyArray(self._body_com_ang_acc_w) + return self._body_com_ang_acc_w_ta @property - def body_com_quat_b(self) -> wp.array: + def body_com_quat_b(self) -> ProxyArray: """Orientation (x, y, z, w) of the principal axes of inertia of all of the bodies in their respective link frames. Shape is (num_instances, 1), dtype = wp.quatf. In torch this resolves to (num_instances, 1, 4). This quantity is the orientation of the principal axes of inertia relative to its body's link frame. """ - return self._get_quat_from_transform(self._body_com_quat_b, self.body_com_pose_b) + self._body_com_quat_b = self._get_quat_from_transform(self._body_com_quat_b, self.body_com_pose_b.warp) + if self._body_com_quat_b_ta is None: + self._body_com_quat_b_ta = ProxyArray(self._body_com_quat_b) + return self._body_com_quat_b_ta def _create_simulation_bindings(self) -> None: """Create simulation bindings for the root data. @@ -751,6 +831,11 @@ def _create_simulation_bindings(self) -> None: :, 0 ] + # Re-pin ProxyArray wrappers to the newly created sim bindings. + # On first init, _create_buffers() handles this after all buffers exist. + if hasattr(self, "_root_link_pose_w_ta"): + self._pin_proxy_arrays() + def _create_buffers(self) -> None: """Create buffers for the root data.""" super()._create_buffers() @@ -833,6 +918,108 @@ def _create_buffers(self) -> None: self._body_com_lin_acc_w = None self._body_com_ang_acc_w = None + # Pin all ProxyArray wrappers to current buffers. + self._pin_proxy_arrays() + + def _pin_proxy_arrays(self) -> None: + """Create or rebind all pinned ProxyArray wrappers. + + Called from :meth:`_create_buffers` on first initialization and from + :meth:`_create_simulation_bindings` after a full simulation reset when + the solver recreates its internal arrays. + """ + is_rebind = hasattr(self, "_root_link_pose_w_ta") + + if is_rebind: + # Rebind sim-bound ProxyArrays to new solver arrays + self._root_link_pose_w_ta = ProxyArray(self._sim_bind_root_link_pose_w) + self._root_com_vel_w_ta = ProxyArray(self._sim_bind_root_com_vel_w) + self._body_link_pose_w_ta = ProxyArray(self._sim_bind_body_link_pose_w) + self._body_com_vel_w_ta = ProxyArray(self._sim_bind_body_com_vel_w) + self._body_mass_ta = ProxyArray(self._sim_bind_body_mass) + self._body_inertia_ta = ProxyArray(self._sim_bind_body_inertia) + self._body_com_pos_b_ta = ProxyArray(self._sim_bind_body_com_pos_b) + else: + # First-time creation: pin ProxyArrays to current buffers + # Category 1: sim-bound and pre-allocated buffers + # Newton wp.array pointers are stable, so a ProxyArray wrapping them is valid forever. + self._root_link_pose_w_ta = ProxyArray(self._sim_bind_root_link_pose_w) + self._root_com_vel_w_ta = ProxyArray(self._sim_bind_root_com_vel_w) + self._body_link_pose_w_ta = ProxyArray(self._sim_bind_body_link_pose_w) + self._body_com_vel_w_ta = ProxyArray(self._sim_bind_body_com_vel_w) + self._default_root_pose_ta = ProxyArray(self._default_root_pose) + self._default_root_vel_ta = ProxyArray(self._default_root_vel) + self._body_mass_ta = ProxyArray(self._sim_bind_body_mass) + self._body_inertia_ta = ProxyArray(self._sim_bind_body_inertia) + self._body_com_pos_b_ta = ProxyArray(self._sim_bind_body_com_pos_b) + + # Category 2: TimestampedBuffer properties + self._root_link_vel_w_ta = ProxyArray(self._root_link_vel_w.data) + self._root_com_pose_w_ta = ProxyArray(self._root_com_pose_w.data) + self._body_com_acc_w_ta = ProxyArray(self._body_com_acc_w.data) + self._body_com_pose_b_ta = ProxyArray(self._body_com_pose_b.data) + self._projected_gravity_b_ta = ProxyArray(self._projected_gravity_b.data) + self._heading_w_ta = ProxyArray(self._heading_w.data) + + # -- deprecated state properties (lazy); type annotations declared once here + self._root_link_lin_vel_b_ta: ProxyArray | None = None + self._root_link_ang_vel_b_ta: ProxyArray | None = None + self._root_com_lin_vel_b_ta: ProxyArray | None = None + self._root_com_ang_vel_b_ta: ProxyArray | None = None + self._root_state_w_ta: ProxyArray | None = None + self._root_link_state_w_ta: ProxyArray | None = None + self._root_com_state_w_ta: ProxyArray | None = None + self._default_root_state_ta: ProxyArray | None = None + self._body_state_w_ta: ProxyArray | None = None + self._body_link_state_w_ta: ProxyArray | None = None + self._body_com_state_w_ta: ProxyArray | None = None + + # Invalidate lazy sliced ProxyArrays AND their backing wp.arrays so they are + # re-created from fresh data on next access. On first init the backing fields + # are already None (set by _create_buffers), so the assignments below are + # harmless no-ops. On rebind they reset stale pointers into freed transform + # memory after a sim reset. + self._root_link_pos_w_ta: ProxyArray | None = None + self._root_link_pos_w = None + self._root_link_quat_w_ta: ProxyArray | None = None + self._root_link_quat_w = None + self._root_link_lin_vel_w_ta: ProxyArray | None = None + self._root_link_lin_vel_w = None + self._root_link_ang_vel_w_ta: ProxyArray | None = None + self._root_link_ang_vel_w = None + self._root_com_pos_w_ta: ProxyArray | None = None + self._root_com_pos_w = None + self._root_com_quat_w_ta: ProxyArray | None = None + self._root_com_quat_w = None + self._root_com_lin_vel_w_ta: ProxyArray | None = None + self._root_com_lin_vel_w = None + self._root_com_ang_vel_w_ta: ProxyArray | None = None + self._root_com_ang_vel_w = None + self._body_link_pos_w_ta: ProxyArray | None = None + self._body_link_pos_w = None + self._body_link_quat_w_ta: ProxyArray | None = None + self._body_link_quat_w = None + self._body_link_vel_w_ta: ProxyArray | None = None + self._body_link_lin_vel_w_ta: ProxyArray | None = None + self._body_link_lin_vel_w = None + self._body_link_ang_vel_w_ta: ProxyArray | None = None + self._body_link_ang_vel_w = None + self._body_com_pose_w_ta: ProxyArray | None = None + self._body_com_pos_w_ta: ProxyArray | None = None + self._body_com_pos_w = None + self._body_com_quat_w_ta: ProxyArray | None = None + self._body_com_quat_w = None + self._body_com_lin_vel_w_ta: ProxyArray | None = None + self._body_com_lin_vel_w = None + self._body_com_ang_vel_w_ta: ProxyArray | None = None + self._body_com_ang_vel_w = None + self._body_com_lin_acc_w_ta: ProxyArray | None = None + self._body_com_lin_acc_w = None + self._body_com_ang_acc_w_ta: ProxyArray | None = None + self._body_com_ang_acc_w = None + self._body_com_quat_b_ta: ProxyArray | None = None + self._body_com_quat_b = None + """ Internal helpers. """ @@ -1037,7 +1224,7 @@ def _get_bottom_from_spatial_vector(self, source: wp.array | None, spatial_vecto """ @property - def root_state_w(self) -> wp.array: + def root_state_w(self) -> ProxyArray: """Deprecated, same as :attr:`root_link_pose_w` and :attr:`root_com_vel_w`.""" warnings.warn( "The `root_state_w` property will be deprecated in IsaacLab 4.0. Please use `root_link_pose_w` and " @@ -1049,13 +1236,14 @@ def root_state_w(self) -> wp.array: self._root_state_w = TimestampedBuffer( shape=(self._num_instances,), dtype=shared_kernels.vec13f, device=self.device ) + self._root_state_w_ta = ProxyArray(self._root_state_w.data) if self._root_state_w.timestamp < self._sim_timestamp: wp.launch( shared_kernels.concat_root_pose_and_vel_to_state, dim=self._num_instances, inputs=[ - self.root_link_pose_w, - self.root_com_vel_w, + self.root_link_pose_w.warp, + self.root_com_vel_w.warp, ], outputs=[ self._root_state_w.data, @@ -1064,10 +1252,10 @@ def root_state_w(self) -> wp.array: ) self._root_state_w.timestamp = self._sim_timestamp - return self._root_state_w.data + return self._root_state_w_ta @property - def root_link_state_w(self) -> wp.array: + def root_link_state_w(self) -> ProxyArray: """Deprecated, same as :attr:`root_link_pose_w` and :attr:`root_link_vel_w`.""" warnings.warn( "The `root_link_state_w` property will be deprecated in IsaacLab 4.0. Please use `root_link_pose_w` and " @@ -1079,13 +1267,14 @@ def root_link_state_w(self) -> wp.array: self._root_link_state_w = TimestampedBuffer( shape=(self._num_instances,), dtype=shared_kernels.vec13f, device=self.device ) + self._root_link_state_w_ta = ProxyArray(self._root_link_state_w.data) if self._root_link_state_w.timestamp < self._sim_timestamp: wp.launch( shared_kernels.concat_root_pose_and_vel_to_state, dim=self._num_instances, inputs=[ - self.root_link_pose_w, - self.root_link_vel_w, + self.root_link_pose_w.warp, + self.root_link_vel_w.warp, ], outputs=[ self._root_link_state_w.data, @@ -1094,10 +1283,10 @@ def root_link_state_w(self) -> wp.array: ) self._root_link_state_w.timestamp = self._sim_timestamp - return self._root_link_state_w.data + return self._root_link_state_w_ta @property - def root_com_state_w(self) -> wp.array: + def root_com_state_w(self) -> ProxyArray: """Deprecated, same as :attr:`root_com_pose_w` and :attr:`root_com_vel_w`.""" warnings.warn( "The `root_com_state_w` property will be deprecated in IsaacLab 4.0. Please use `root_com_pose_w` and " @@ -1109,13 +1298,14 @@ def root_com_state_w(self) -> wp.array: self._root_com_state_w = TimestampedBuffer( shape=(self._num_instances,), dtype=shared_kernels.vec13f, device=self.device ) + self._root_com_state_w_ta = ProxyArray(self._root_com_state_w.data) if self._root_com_state_w.timestamp < self._sim_timestamp: wp.launch( shared_kernels.concat_root_pose_and_vel_to_state, dim=self._num_instances, inputs=[ - self.root_com_pose_w, - self.root_com_vel_w, + self.root_com_pose_w.warp, + self.root_com_vel_w.warp, ], outputs=[ self._root_com_state_w.data, @@ -1124,10 +1314,10 @@ def root_com_state_w(self) -> wp.array: ) self._root_com_state_w.timestamp = self._sim_timestamp - return self._root_com_state_w.data + return self._root_com_state_w_ta @property - def default_root_state(self) -> wp.array: + def default_root_state(self) -> ProxyArray: """Default root state ``[pos, quat, lin_vel, ang_vel]`` in local environment frame. The position and quaternion are of the rigid body's actor frame. Meanwhile, the linear and angular velocities @@ -1141,6 +1331,7 @@ def default_root_state(self) -> wp.array: ) if self._default_root_state is None: self._default_root_state = wp.zeros((self._num_instances), dtype=shared_kernels.vec13f, device=self.device) + self._default_root_state_ta = ProxyArray(self._default_root_state) wp.launch( shared_kernels.concat_root_pose_and_vel_to_state, dim=self._num_instances, @@ -1153,10 +1344,10 @@ def default_root_state(self) -> wp.array: ], device=self.device, ) - return self._default_root_state + return self._default_root_state_ta @property - def body_state_w(self) -> wp.array: + def body_state_w(self) -> ProxyArray: """Deprecated, same as :attr:`body_link_pose_w` and :attr:`body_com_vel_w`.""" warnings.warn( "The `body_state_w` property will be deprecated in IsaacLab 4.0. Please use `body_link_pose_w` and " @@ -1169,13 +1360,14 @@ def body_state_w(self) -> wp.array: self._root_state_w = TimestampedBuffer( shape=(self._num_instances,), dtype=shared_kernels.vec13f, device=self.device ) + self._root_state_w_ta = ProxyArray(self._root_state_w.data) if self._root_state_w.timestamp < self._sim_timestamp: wp.launch( shared_kernels.concat_root_pose_and_vel_to_state, dim=self._num_instances, inputs=[ - self.root_link_pose_w, - self.root_com_vel_w, + self.root_link_pose_w.warp, + self.root_com_vel_w.warp, ], outputs=[ self._root_state_w.data, @@ -1183,10 +1375,12 @@ def body_state_w(self) -> wp.array: device=self.device, ) self._root_state_w.timestamp = self._sim_timestamp - return self._root_state_w.data.reshape((self._num_instances, 1)) + if self._body_state_w_ta is None: + self._body_state_w_ta = ProxyArray(self._root_state_w.data.reshape((self._num_instances, 1))) + return self._body_state_w_ta @property - def body_link_state_w(self) -> wp.array: + def body_link_state_w(self) -> ProxyArray: """Deprecated, same as :attr:`body_link_pose_w` and :attr:`body_link_vel_w`.""" warnings.warn( "The `body_link_state_w` property will be deprecated in IsaacLab 4.0. Please use `body_link_pose_w` and " @@ -1199,13 +1393,14 @@ def body_link_state_w(self) -> wp.array: self._root_link_state_w = TimestampedBuffer( shape=(self._num_instances,), dtype=shared_kernels.vec13f, device=self.device ) + self._root_link_state_w_ta = ProxyArray(self._root_link_state_w.data) if self._root_link_state_w.timestamp < self._sim_timestamp: wp.launch( shared_kernels.concat_root_pose_and_vel_to_state, dim=self._num_instances, inputs=[ - self.root_link_pose_w, - self.root_link_vel_w, + self.root_link_pose_w.warp, + self.root_link_vel_w.warp, ], outputs=[ self._root_link_state_w.data, @@ -1213,10 +1408,12 @@ def body_link_state_w(self) -> wp.array: device=self.device, ) self._root_link_state_w.timestamp = self._sim_timestamp - return self._root_link_state_w.data.reshape((self._num_instances, 1)) + if self._body_link_state_w_ta is None: + self._body_link_state_w_ta = ProxyArray(self._root_link_state_w.data.reshape((self._num_instances, 1))) + return self._body_link_state_w_ta @property - def body_com_state_w(self) -> wp.array: + def body_com_state_w(self) -> ProxyArray: """Deprecated, same as :attr:`body_com_pose_w` and :attr:`body_com_vel_w`.""" warnings.warn( "The `body_com_state_w` property will be deprecated in IsaacLab 4.0. Please use `body_com_pose_w` and " @@ -1228,13 +1425,14 @@ def body_com_state_w(self) -> wp.array: self._root_com_state_w = TimestampedBuffer( shape=(self._num_instances,), dtype=shared_kernels.vec13f, device=self.device ) + self._root_com_state_w_ta = ProxyArray(self._root_com_state_w.data) if self._root_com_state_w.timestamp < self._sim_timestamp: wp.launch( shared_kernels.concat_root_pose_and_vel_to_state, dim=self._num_instances, inputs=[ - self.root_com_pose_w, - self.root_com_vel_w, + self.root_com_pose_w.warp, + self.root_com_vel_w.warp, ], outputs=[ self._root_com_state_w.data, @@ -1242,4 +1440,6 @@ def body_com_state_w(self) -> wp.array: device=self.device, ) self._root_com_state_w.timestamp = self._sim_timestamp - return self._root_com_state_w.data.reshape((self._num_instances, 1)) + if self._body_com_state_w_ta is None: + self._body_com_state_w_ta = ProxyArray(self._root_com_state_w.data.reshape((self._num_instances, 1))) + return self._body_com_state_w_ta diff --git a/source/isaaclab_newton/isaaclab_newton/assets/rigid_object_collection/rigid_object_collection.py b/source/isaaclab_newton/isaaclab_newton/assets/rigid_object_collection/rigid_object_collection.py index 52216290d329..0dbe9e981cc0 100644 --- a/source/isaaclab_newton/isaaclab_newton/assets/rigid_object_collection/rigid_object_collection.py +++ b/source/isaaclab_newton/isaaclab_newton/assets/rigid_object_collection/rigid_object_collection.py @@ -21,6 +21,7 @@ import isaaclab.sim as sim_utils import isaaclab.utils.string as string_utils from isaaclab.assets.rigid_object_collection.base_rigid_object_collection import BaseRigidObjectCollection +from isaaclab.physics import PhysicsEvent from isaaclab.utils.wrench_composer import WrenchComposer from isaaclab_newton.assets import kernels as shared_kernels @@ -1102,6 +1103,13 @@ def _initialize_impl(self): # container for data access self._data = RigidObjectCollectionData(self._root_view, self.num_bodies, self.device) + # Register callback to rebind simulation data after a full reset (model/state recreation). + self._physics_ready_handle = SimulationManager.register_callback( + lambda _: self._data._create_simulation_bindings(), + PhysicsEvent.PHYSICS_READY, + name=f"rigid_object_collection_rebind_{self.cfg.rigid_objects}", + ) + # create buffers self._create_buffers() # process configuration diff --git a/source/isaaclab_newton/isaaclab_newton/assets/rigid_object_collection/rigid_object_collection_data.py b/source/isaaclab_newton/isaaclab_newton/assets/rigid_object_collection/rigid_object_collection_data.py index 3fbc1801322f..9fe3ec125de0 100644 --- a/source/isaaclab_newton/isaaclab_newton/assets/rigid_object_collection/rigid_object_collection_data.py +++ b/source/isaaclab_newton/isaaclab_newton/assets/rigid_object_collection/rigid_object_collection_data.py @@ -15,6 +15,7 @@ from isaaclab.assets.rigid_object_collection.base_rigid_object_collection_data import BaseRigidObjectCollectionData from isaaclab.utils.buffers import TimestampedBufferWarp as TimestampedBuffer from isaaclab.utils.math import normalize +from isaaclab.utils.warp import ProxyArray from isaaclab_newton.assets import kernels as shared_kernels from isaaclab_newton.physics import NewtonManager as SimulationManager @@ -75,10 +76,14 @@ def __init__(self, root_view: ArticulationView, num_bodies: int, device: str): gravity_dir = normalize(gravity_dir.unsqueeze(0)).squeeze(0) # Initialize constants - self.GRAVITY_VEC_W = wp.from_torch(gravity_dir.repeat(self.num_instances, self.num_bodies, 1), dtype=wp.vec3f) - self.FORWARD_VEC_B = wp.from_torch( - torch.tensor((1.0, 0.0, 0.0), device=self.device).repeat(self.num_instances, self.num_bodies, 1), - dtype=wp.vec3f, + self.GRAVITY_VEC_W = ProxyArray( + wp.from_torch(gravity_dir.repeat(self.num_instances, self.num_bodies, 1), dtype=wp.vec3f) + ) + self.FORWARD_VEC_B = ProxyArray( + wp.from_torch( + torch.tensor((1.0, 0.0, 0.0), device=self.device).repeat(self.num_instances, self.num_bodies, 1), + dtype=wp.vec3f, + ) ) self._create_simulation_bindings() @@ -130,14 +135,14 @@ def update(self, dt: float) -> None: """ @property - def default_body_pose(self) -> wp.array: + def default_body_pose(self) -> ProxyArray: """Default body pose ``[pos, quat]`` in local environment frame. The position and quaternion are of the rigid body's actor frame. Shape is (num_instances, num_bodies), dtype = wp.transformf. In torch this resolves to (num_instances, num_bodies, 7). """ - return self._default_body_pose + return self._default_body_pose_ta @default_body_pose.setter def default_body_pose(self, value: wp.array) -> None: @@ -154,14 +159,14 @@ def default_body_pose(self, value: wp.array) -> None: self._default_body_pose.assign(value) @property - def default_body_vel(self) -> wp.array: + def default_body_vel(self) -> ProxyArray: """Default body velocity ``[lin_vel, ang_vel]`` in local environment frame. The linear and angular velocities are of the rigid body's center of mass frame. Shape is (num_instances, num_bodies), dtype = wp.spatial_vectorf. In torch this resolves to (num_instances, num_bodies, 6). """ - return self._default_body_vel + return self._default_body_vel_ta @default_body_vel.setter def default_body_vel(self, value: wp.array) -> None: @@ -182,7 +187,7 @@ def default_body_vel(self, value: wp.array) -> None: """ @property - def body_link_pose_w(self) -> wp.array: + def body_link_pose_w(self) -> ProxyArray: """Body link pose ``[pos, quat]`` in simulation world frame. Shape is (num_instances, num_bodies), dtype = wp.transformf. In torch this resolves to @@ -190,10 +195,10 @@ def body_link_pose_w(self) -> wp.array: 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. """ - return self._sim_bind_body_link_pose_w + return self._body_link_pose_w_ta @property - def body_link_vel_w(self) -> wp.array: + def body_link_vel_w(self) -> ProxyArray: """Body link velocity ``[lin_vel, ang_vel]`` in simulation world frame. Shape is (num_instances, num_bodies), dtype = wp.spatial_vectorf. In torch this resolves to @@ -206,9 +211,9 @@ def body_link_vel_w(self) -> wp.array: shared_kernels.get_body_link_vel_from_body_com_vel, dim=(self.num_instances, self.num_bodies), inputs=[ - self.body_com_vel_w, - self.body_link_pose_w, - self.body_com_pos_b, + self.body_com_vel_w.warp, + self.body_link_pose_w.warp, + self.body_com_pos_b.warp, ], outputs=[ self._body_link_vel_w.data, @@ -217,10 +222,10 @@ def body_link_vel_w(self) -> wp.array: ) self._body_link_vel_w.timestamp = self._sim_timestamp - return self._body_link_vel_w.data + return self._body_link_vel_w_ta @property - def body_com_pose_w(self) -> wp.array: + def body_com_pose_w(self) -> ProxyArray: """Body center of mass pose ``[pos, quat]`` in simulation world frame. Shape is (num_instances, num_bodies), dtype = wp.transformf. In torch this resolves to @@ -233,8 +238,8 @@ def body_com_pose_w(self) -> wp.array: shared_kernels.get_body_com_pose_from_body_link_pose, dim=(self.num_instances, self.num_bodies), inputs=[ - self.body_link_pose_w, - self.body_com_pos_b, + self.body_link_pose_w.warp, + self.body_com_pos_b.warp, ], outputs=[ self._body_com_pose_w.data, @@ -243,10 +248,10 @@ def body_com_pose_w(self) -> wp.array: ) self._body_com_pose_w.timestamp = self._sim_timestamp - return self._body_com_pose_w.data + return self._body_com_pose_w_ta @property - def body_com_vel_w(self) -> wp.array: + def body_com_vel_w(self) -> ProxyArray: """Body center of mass velocity ``[lin_vel, ang_vel]`` in simulation world frame. Shape is (num_instances, num_bodies), dtype = wp.spatial_vectorf. In torch this resolves to @@ -254,10 +259,10 @@ def body_com_vel_w(self) -> wp.array: This quantity contains the linear and angular velocities of the root rigid body's center of mass frame relative to the world. """ - return self._sim_bind_body_com_vel_w + return self._body_com_vel_w_ta @property - def body_com_acc_w(self) -> wp.array: + def body_com_acc_w(self) -> ProxyArray: """Acceleration of all bodies ``[lin_acc, ang_acc]`` in the simulation world frame. Shape is (num_instances, num_bodies), dtype = wp.spatial_vectorf. In torch this resolves to @@ -270,7 +275,7 @@ def body_com_acc_w(self) -> wp.array: dim=(self.num_instances, self.num_bodies), device=self.device, inputs=[ - self.body_com_vel_w, + self.body_com_vel_w.warp, SimulationManager.get_dt(), self._previous_body_com_vel, ], @@ -279,10 +284,10 @@ def body_com_acc_w(self) -> wp.array: ], ) self._body_com_acc_w.timestamp = self._sim_timestamp - return self._body_com_acc_w.data + return self._body_com_acc_w_ta @property - def body_com_pose_b(self) -> wp.array: + def body_com_pose_b(self) -> ProxyArray: """Center of mass pose ``[pos, quat]`` of all bodies in their respective body's link frames. Shape is (num_instances, num_bodies), dtype = wp.transformf. In torch this resolves to @@ -301,7 +306,7 @@ def body_com_pose_b(self) -> wp.array: shared_kernels.make_dummy_body_com_pose_b, dim=(self.num_instances, self.num_bodies), inputs=[ - self.body_com_pos_b, + self.body_com_pos_b.warp, ], outputs=[ self._body_com_pose_b.data, @@ -309,42 +314,42 @@ def body_com_pose_b(self) -> wp.array: device=self.device, ) self._body_com_pose_b.timestamp = self._sim_timestamp - return self._body_com_pose_b.data + return self._body_com_pose_b_ta @property - def body_com_pos_b(self) -> wp.array: + def body_com_pos_b(self) -> ProxyArray: """Center of mass position of all of the bodies in their respective link frames. Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to (num_instances, num_bodies, 3). This quantity is the center of mass location relative to its body's link frame. """ - return self._sim_bind_body_com_pos_b + return self._body_com_pos_b_ta @property - def body_mass(self) -> wp.array: + def body_mass(self) -> ProxyArray: """Mass of all bodies in the simulation world frame. Shape is (num_instances, num_bodies), dtype = wp.float32. In torch this resolves to (num_instances, num_bodies). """ - return self._sim_bind_body_mass + return self._body_mass_ta @property - def body_inertia(self) -> wp.array: + def body_inertia(self) -> ProxyArray: """Inertia of all bodies in the simulation world frame. Shape is (num_instances, num_bodies, 9), dtype = wp.float32. In torch this resolves to (num_instances, num_bodies, 9). """ - return self._body_inertia + return self._body_inertia_ta """ Derived Properties. """ @property - def projected_gravity_b(self) -> wp.array: + def projected_gravity_b(self) -> ProxyArray: """Projection of the gravity direction on base frame. Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to @@ -354,15 +359,15 @@ def projected_gravity_b(self) -> wp.array: wp.launch( shared_kernels.quat_apply_inverse_2D_kernel, dim=(self.num_instances, self.num_bodies), - inputs=[self.GRAVITY_VEC_W, self.body_link_quat_w], + inputs=[self.GRAVITY_VEC_W.warp, self.body_link_quat_w.warp], outputs=[self._projected_gravity_b.data], device=self.device, ) self._projected_gravity_b.timestamp = self._sim_timestamp - return self._projected_gravity_b.data + return self._projected_gravity_b_ta @property - def heading_w(self) -> wp.array: + def heading_w(self) -> ProxyArray: """Yaw heading of the base frame (in radians). Shape is (num_instances, num_bodies), dtype = wp.float32. In torch this resolves to (num_instances, num_bodies). @@ -375,15 +380,15 @@ def heading_w(self) -> wp.array: wp.launch( shared_kernels.body_heading_w, dim=(self.num_instances, self.num_bodies), - inputs=[self.FORWARD_VEC_B, self.body_link_quat_w], + inputs=[self.FORWARD_VEC_B.warp, self.body_link_quat_w.warp], outputs=[self._heading_w.data], device=self.device, ) self._heading_w.timestamp = self._sim_timestamp - return self._heading_w.data + return self._heading_w_ta @property - def body_link_lin_vel_b(self) -> wp.array: + def body_link_lin_vel_b(self) -> ProxyArray: """Root link linear velocity in base frame. Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to @@ -395,15 +400,15 @@ def body_link_lin_vel_b(self) -> wp.array: wp.launch( shared_kernels.quat_apply_inverse_2D_kernel, dim=(self.num_instances, self.num_bodies), - inputs=[self.body_link_lin_vel_w, self.body_link_quat_w], + inputs=[self.body_link_lin_vel_w.warp, self.body_link_quat_w.warp], outputs=[self._body_link_lin_vel_b.data], device=self.device, ) self._body_link_lin_vel_b.timestamp = self._sim_timestamp - return self._body_link_lin_vel_b.data + return self._body_link_lin_vel_b_ta @property - def body_link_ang_vel_b(self) -> wp.array: + def body_link_ang_vel_b(self) -> ProxyArray: """Root link angular velocity in base frame. Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to @@ -415,15 +420,15 @@ def body_link_ang_vel_b(self) -> wp.array: wp.launch( shared_kernels.quat_apply_inverse_2D_kernel, dim=(self.num_instances, self.num_bodies), - inputs=[self.body_link_ang_vel_w, self.body_link_quat_w], + inputs=[self.body_link_ang_vel_w.warp, self.body_link_quat_w.warp], outputs=[self._body_link_ang_vel_b.data], device=self.device, ) self._body_link_ang_vel_b.timestamp = self._sim_timestamp - return self._body_link_ang_vel_b.data + return self._body_link_ang_vel_b_ta @property - def body_com_lin_vel_b(self) -> wp.array: + def body_com_lin_vel_b(self) -> ProxyArray: """Root center of mass linear velocity in base frame. Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to @@ -435,15 +440,15 @@ def body_com_lin_vel_b(self) -> wp.array: wp.launch( shared_kernels.quat_apply_inverse_2D_kernel, dim=(self.num_instances, self.num_bodies), - inputs=[self.body_com_lin_vel_w, self.body_link_quat_w], + inputs=[self.body_com_lin_vel_w.warp, self.body_link_quat_w.warp], outputs=[self._body_com_lin_vel_b.data], device=self.device, ) self._body_com_lin_vel_b.timestamp = self._sim_timestamp - return self._body_com_lin_vel_b.data + return self._body_com_lin_vel_b_ta @property - def body_com_ang_vel_b(self) -> wp.array: + def body_com_ang_vel_b(self) -> ProxyArray: """Root center of mass angular velocity in base frame. Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to @@ -455,119 +460,139 @@ def body_com_ang_vel_b(self) -> wp.array: wp.launch( shared_kernels.quat_apply_inverse_2D_kernel, dim=(self.num_instances, self.num_bodies), - inputs=[self.body_com_ang_vel_w, self.body_link_quat_w], + inputs=[self.body_com_ang_vel_w.warp, self.body_link_quat_w.warp], outputs=[self._body_com_ang_vel_b.data], device=self.device, ) self._body_com_ang_vel_b.timestamp = self._sim_timestamp - return self._body_com_ang_vel_b.data + return self._body_com_ang_vel_b_ta """ Sliced properties. """ @property - def body_link_pos_w(self) -> wp.array: + def body_link_pos_w(self) -> ProxyArray: """Positions of all bodies in simulation world frame. Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to (num_instances, num_bodies, 3). This quantity is the position of the rigid bodies' actor frame relative to the world. """ - return self._get_pos_from_transform(self.body_link_pose_w) + if self._body_link_pos_w_ta is None: + self._body_link_pos_w_ta = ProxyArray(self._get_pos_from_transform(self.body_link_pose_w.warp)) + return self._body_link_pos_w_ta @property - def body_link_quat_w(self) -> wp.array: + def body_link_quat_w(self) -> ProxyArray: """Orientation (x, y, z, w) of all bodies in simulation world frame. Shape is (num_instances, num_bodies), dtype = wp.quatf. In torch this resolves to (num_instances, num_bodies, 4). This quantity is the orientation of the rigid bodies' actor frame relative to the world. """ - return self._get_quat_from_transform(self.body_link_pose_w) + if self._body_link_quat_w_ta is None: + self._body_link_quat_w_ta = ProxyArray(self._get_quat_from_transform(self.body_link_pose_w.warp)) + return self._body_link_quat_w_ta @property - def body_link_lin_vel_w(self) -> wp.array: + def body_link_lin_vel_w(self) -> ProxyArray: """Linear velocity of all bodies in simulation world frame. Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to (num_instances, num_bodies, 3). This quantity is the linear velocity of the rigid bodies' actor frame relative to the world. """ - return self._get_lin_vel_from_spatial_vector(self.body_link_vel_w) + if self._body_link_lin_vel_w_ta is None: + self._body_link_lin_vel_w_ta = ProxyArray(self._get_lin_vel_from_spatial_vector(self.body_link_vel_w.warp)) + return self._body_link_lin_vel_w_ta @property - def body_link_ang_vel_w(self) -> wp.array: + def body_link_ang_vel_w(self) -> ProxyArray: """Angular velocity of all bodies in simulation world frame. Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to (num_instances, num_bodies, 3). This quantity is the angular velocity of the rigid bodies' actor frame relative to the world. """ - return self._get_ang_vel_from_spatial_vector(self.body_link_vel_w) + if self._body_link_ang_vel_w_ta is None: + self._body_link_ang_vel_w_ta = ProxyArray(self._get_ang_vel_from_spatial_vector(self.body_link_vel_w.warp)) + return self._body_link_ang_vel_w_ta @property - def body_com_pos_w(self) -> wp.array: + def body_com_pos_w(self) -> ProxyArray: """Positions of all bodies' center of mass in simulation world frame. Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to (num_instances, num_bodies, 3). This quantity is the position of the rigid bodies' center of mass frame. """ - return self._get_pos_from_transform(self.body_com_pose_w) + if self._body_com_pos_w_ta is None: + self._body_com_pos_w_ta = ProxyArray(self._get_pos_from_transform(self.body_com_pose_w.warp)) + return self._body_com_pos_w_ta @property - def body_com_quat_w(self) -> wp.array: + def body_com_quat_w(self) -> ProxyArray: """Orientation (x, y, z, w) of the principal axes of inertia of all bodies in simulation world frame. Shape is (num_instances, num_bodies), dtype = wp.quatf. In torch this resolves to (num_instances, num_bodies, 4). This quantity is the orientation of the principal axes of inertia of the rigid bodies. """ - return self._get_quat_from_transform(self.body_com_pose_w) + if self._body_com_quat_w_ta is None: + self._body_com_quat_w_ta = ProxyArray(self._get_quat_from_transform(self.body_com_pose_w.warp)) + return self._body_com_quat_w_ta @property - def body_com_lin_vel_w(self) -> wp.array: + def body_com_lin_vel_w(self) -> ProxyArray: """Linear velocity of all bodies in simulation world frame. Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to (num_instances, num_bodies, 3). This quantity is the linear velocity of the rigid bodies' center of mass frame. """ - return self._get_lin_vel_from_spatial_vector(self.body_com_vel_w) + if self._body_com_lin_vel_w_ta is None: + self._body_com_lin_vel_w_ta = ProxyArray(self._get_lin_vel_from_spatial_vector(self.body_com_vel_w.warp)) + return self._body_com_lin_vel_w_ta @property - def body_com_ang_vel_w(self) -> wp.array: + def body_com_ang_vel_w(self) -> ProxyArray: """Angular velocity of all bodies in simulation world frame. Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to (num_instances, num_bodies, 3). This quantity is the angular velocity of the rigid bodies' center of mass frame. """ - return self._get_ang_vel_from_spatial_vector(self.body_com_vel_w) + if self._body_com_ang_vel_w_ta is None: + self._body_com_ang_vel_w_ta = ProxyArray(self._get_ang_vel_from_spatial_vector(self.body_com_vel_w.warp)) + return self._body_com_ang_vel_w_ta @property - def body_com_lin_acc_w(self) -> wp.array: + def body_com_lin_acc_w(self) -> ProxyArray: """Linear acceleration of all bodies in simulation world frame. Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to (num_instances, num_bodies, 3). This quantity is the linear acceleration of the rigid bodies' center of mass frame. """ - return self._get_lin_vel_from_spatial_vector(self.body_com_acc_w) + if self._body_com_lin_acc_w_ta is None: + self._body_com_lin_acc_w_ta = ProxyArray(self._get_lin_vel_from_spatial_vector(self.body_com_acc_w.warp)) + return self._body_com_lin_acc_w_ta @property - def body_com_ang_acc_w(self) -> wp.array: + def body_com_ang_acc_w(self) -> ProxyArray: """Angular acceleration of all bodies in simulation world frame. Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to (num_instances, num_bodies, 3). This quantity is the angular acceleration of the rigid bodies' center of mass frame. """ - return self._get_ang_vel_from_spatial_vector(self.body_com_acc_w) + if self._body_com_ang_acc_w_ta is None: + self._body_com_ang_acc_w_ta = ProxyArray(self._get_ang_vel_from_spatial_vector(self.body_com_acc_w.warp)) + return self._body_com_ang_acc_w_ta @property - def body_com_quat_b(self) -> wp.array: + def body_com_quat_b(self) -> ProxyArray: """Orientation (x, y, z, w) of the principal axes of inertia of all of the bodies in their respective link frames. @@ -575,7 +600,9 @@ def body_com_quat_b(self) -> wp.array: (num_instances, num_bodies, 4). This quantity is the orientation of the principal axes of inertia relative to its body's link frame. """ - return self._get_quat_from_transform(self.body_com_pose_b) + if self._body_com_quat_b_ta is None: + self._body_com_quat_b_ta = ProxyArray(self._get_quat_from_transform(self.body_com_pose_b.warp)) + return self._body_com_quat_b_ta def _create_simulation_bindings(self) -> None: """Create simulation bindings for the body data. @@ -609,6 +636,11 @@ def _create_simulation_bindings(self) -> None: copy=False, ) + # Re-pin ProxyArray wrappers to the newly created sim bindings. + # On first init, _create_buffers() handles this after all buffers exist. + if hasattr(self, "_body_link_pose_w_ta"): + self._pin_proxy_arrays() + def _create_buffers(self) -> None: """Create buffers for computing and caching derived quantities.""" super()._create_buffers() @@ -654,6 +686,74 @@ def _create_buffers(self) -> None: # -- Initialize history for finite differencing self._previous_body_com_vel = wp.clone(self._sim_bind_body_com_vel_w) + # Pin all ProxyArray wrappers to current buffers. + self._pin_proxy_arrays() + + def _pin_proxy_arrays(self) -> None: + """Create or rebind all pinned ProxyArray wrappers. + + Called from :meth:`_create_buffers` on first initialization and from + :meth:`_create_simulation_bindings` after a full simulation reset when + the solver recreates its internal arrays. + """ + is_rebind = hasattr(self, "_body_link_pose_w_ta") + + if is_rebind: + # Rebind sim-bound ProxyArrays to new solver arrays + self._body_link_pose_w_ta = ProxyArray(self._sim_bind_body_link_pose_w) + self._body_com_vel_w_ta = ProxyArray(self._sim_bind_body_com_vel_w) + self._body_com_pos_b_ta = ProxyArray(self._sim_bind_body_com_pos_b) + self._body_mass_ta = ProxyArray(self._sim_bind_body_mass) + self._body_inertia_ta = ProxyArray(self._body_inertia) + else: + # First-time creation: pin ProxyArrays to current buffers + # Category 1: sim-bound and pre-allocated buffers + # Newton wp.array pointers are stable, so a ProxyArray wrapping them is valid forever. + self._body_link_pose_w_ta = ProxyArray(self._sim_bind_body_link_pose_w) + self._body_com_vel_w_ta = ProxyArray(self._sim_bind_body_com_vel_w) + self._body_com_pos_b_ta = ProxyArray(self._sim_bind_body_com_pos_b) + self._body_mass_ta = ProxyArray(self._sim_bind_body_mass) + self._body_inertia_ta = ProxyArray(self._body_inertia) + self._default_body_pose_ta = ProxyArray(self._default_body_pose) + self._default_body_vel_ta = ProxyArray(self._default_body_vel) + + # Category 2: TimestampedBuffer properties + self._body_link_vel_w_ta = ProxyArray(self._body_link_vel_w.data) + self._body_com_pose_w_ta = ProxyArray(self._body_com_pose_w.data) + self._body_com_acc_w_ta = ProxyArray(self._body_com_acc_w.data) + self._body_com_pose_b_ta = ProxyArray(self._body_com_pose_b.data) + self._projected_gravity_b_ta = ProxyArray(self._projected_gravity_b.data) + self._heading_w_ta = ProxyArray(self._heading_w.data) + self._body_link_lin_vel_b_ta = ProxyArray(self._body_link_lin_vel_b.data) + self._body_link_ang_vel_b_ta = ProxyArray(self._body_link_ang_vel_b.data) + self._body_com_lin_vel_b_ta = ProxyArray(self._body_com_lin_vel_b.data) + self._body_com_ang_vel_b_ta = ProxyArray(self._body_com_ang_vel_b.data) + self._body_state_w_ta = ProxyArray(self._body_state_w.data) + self._body_link_state_w_ta = ProxyArray(self._body_link_state_w.data) + self._body_com_state_w_ta = ProxyArray(self._body_com_state_w.data) + + # -- deprecated state properties (lazy); type annotation declared once here + self._default_body_state_ta: ProxyArray | None = None + + # Invalidate lazy sliced ProxyArrays AND their backing wp.arrays so they are + # re-created from fresh data on next access. On first init the backing fields + # are already None (set by _create_buffers), so the assignments below are + # harmless no-ops. On rebind they reset stale pointers into freed transform + # memory after a sim reset. + self._body_link_pos_w_ta: ProxyArray | None = None + self._body_link_quat_w_ta: ProxyArray | None = None + self._body_link_lin_vel_w_ta: ProxyArray | None = None + self._body_link_ang_vel_w_ta: ProxyArray | None = None + self._body_com_pos_w_ta: ProxyArray | None = None + self._body_com_quat_w_ta: ProxyArray | None = None + self._body_com_lin_vel_w_ta: ProxyArray | None = None + self._body_com_ang_vel_w_ta: ProxyArray | None = None + self._body_com_lin_acc_w_ta: ProxyArray | None = None + self._body_com_ang_acc_w_ta: ProxyArray | None = None + self._body_com_quat_b_ta: ProxyArray | None = None + self._default_body_state_ta: ProxyArray | None = None + self._default_body_state = None + """ Helpers. """ @@ -703,7 +803,7 @@ def _get_ang_vel_from_spatial_vector(self, spatial_vector: wp.array) -> wp.array """ @property - def default_body_state(self) -> wp.array: + def default_body_state(self) -> ProxyArray: """Default root state ``[pos, quat, lin_vel, ang_vel]`` in local environment frame. The position and quaternion are of the rigid body's actor frame. Meanwhile, the linear and angular velocities @@ -719,6 +819,7 @@ def default_body_state(self) -> wp.array: self._default_body_state = wp.zeros( (self.num_instances, self.num_bodies), dtype=shared_kernels.vec13f, device=self.device ) + self._default_body_state_ta = ProxyArray(self._default_body_state) wp.launch( shared_kernels.concat_body_pose_and_vel_to_state, dim=(self.num_instances, self.num_bodies), @@ -731,10 +832,10 @@ def default_body_state(self) -> wp.array: ], device=self.device, ) - return self._default_body_state + return self._default_body_state_ta @property - def body_state_w(self) -> wp.array: + def body_state_w(self) -> ProxyArray: """Deprecated, same as :attr:`body_link_pose_w` and :attr:`body_com_vel_w`.""" warnings.warn( "The `body_state_w` property will be deprecated in IsaacLab 4.0. Please use `body_link_pose_w` and " @@ -747,8 +848,8 @@ def body_state_w(self) -> wp.array: shared_kernels.concat_body_pose_and_vel_to_state, dim=(self.num_instances, self.num_bodies), inputs=[ - self.body_link_pose_w, - self.body_com_vel_w, + self.body_link_pose_w.warp, + self.body_com_vel_w.warp, ], outputs=[ self._body_state_w.data, @@ -757,10 +858,10 @@ def body_state_w(self) -> wp.array: ) self._body_state_w.timestamp = self._sim_timestamp - return self._body_state_w.data + return self._body_state_w_ta @property - def body_link_state_w(self) -> wp.array: + def body_link_state_w(self) -> ProxyArray: """Deprecated, same as :attr:`body_link_pose_w` and :attr:`body_link_vel_w`.""" warnings.warn( "The `body_link_state_w` property will be deprecated in IsaacLab 4.0. Please use `body_link_pose_w` and " @@ -773,8 +874,8 @@ def body_link_state_w(self) -> wp.array: shared_kernels.concat_body_pose_and_vel_to_state, dim=(self.num_instances, self.num_bodies), inputs=[ - self.body_link_pose_w, - self.body_link_vel_w, + self.body_link_pose_w.warp, + self.body_link_vel_w.warp, ], outputs=[ self._body_link_state_w.data, @@ -783,10 +884,10 @@ def body_link_state_w(self) -> wp.array: ) self._body_link_state_w.timestamp = self._sim_timestamp - return self._body_link_state_w.data + return self._body_link_state_w_ta @property - def body_com_state_w(self) -> wp.array: + def body_com_state_w(self) -> ProxyArray: """Deprecated, same as :attr:`body_com_pose_w` and :attr:`body_com_vel_w`.""" warnings.warn( "The `body_com_state_w` property will be deprecated in IsaacLab 4.0. Please use `body_com_pose_w` and " @@ -799,8 +900,8 @@ def body_com_state_w(self) -> wp.array: shared_kernels.concat_body_pose_and_vel_to_state, dim=(self.num_instances, self.num_bodies), inputs=[ - self.body_com_pose_w, - self.body_com_vel_w, + self.body_com_pose_w.warp, + self.body_com_vel_w.warp, ], outputs=[ self._body_com_state_w.data, @@ -809,4 +910,4 @@ def body_com_state_w(self) -> wp.array: ) self._body_com_state_w.timestamp = self._sim_timestamp - return self._body_com_state_w.data + return self._body_com_state_w_ta diff --git a/source/isaaclab_newton/isaaclab_newton/renderers/newton_warp_renderer.py b/source/isaaclab_newton/isaaclab_newton/renderers/newton_warp_renderer.py index 22f8acb47a93..11f1a833839e 100644 --- a/source/isaaclab_newton/isaaclab_newton/renderers/newton_warp_renderer.py +++ b/source/isaaclab_newton/isaaclab_newton/renderers/newton_warp_renderer.py @@ -117,13 +117,13 @@ def update(self, positions: torch.Tensor, orientations: torch.Tensor, intrinsics ) def _from_torch(self, tensor: torch.Tensor, dtype) -> wp.array: - torch_array = wp.from_torch(tensor) + proxy_array = wp.from_torch(tensor) if tensor.is_contiguous(): return wp.array( - ptr=torch_array.ptr, + ptr=proxy_array.ptr, dtype=dtype, shape=(self.newton_sensor.model.world_count, self.num_cameras, self.height, self.width), - device=torch_array.device, + device=proxy_array.device, copy=False, ) @@ -131,7 +131,7 @@ def _from_torch(self, tensor: torch.Tensor, dtype) -> wp.array: return wp.zeros( (self.newton_sensor.model.world_count, self.num_cameras, self.height, self.width), dtype=dtype, - device=torch_array.device, + device=proxy_array.device, ) @wp.kernel diff --git a/source/isaaclab_newton/isaaclab_newton/sensors/contact_sensor/contact_sensor.py b/source/isaaclab_newton/isaaclab_newton/sensors/contact_sensor/contact_sensor.py index 5517b845438b..02850e9cc903 100644 --- a/source/isaaclab_newton/isaaclab_newton/sensors/contact_sensor/contact_sensor.py +++ b/source/isaaclab_newton/isaaclab_newton/sensors/contact_sensor/contact_sensor.py @@ -20,6 +20,7 @@ import isaaclab.utils.string as string_utils from isaaclab.sensors.contact_sensor.base_contact_sensor import BaseContactSensor +from isaaclab.utils.warp import ProxyArray from isaaclab_newton.physics import NewtonManager @@ -207,7 +208,7 @@ def find_sensors(self, name_keys: str | Sequence[str], preserve_order: bool = Fa ) return string_utils.resolve_matching_names(name_keys, sensor_names, preserve_order) - def compute_first_contact(self, dt: float, abs_tol: float = 1.0e-8) -> wp.array: + def compute_first_contact(self, dt: float, abs_tol: float = 1.0e-8) -> ProxyArray: """Checks if sensors that have established contact within the last :attr:`dt` seconds. This function checks if the sensors have established contact within the last :attr:`dt` seconds @@ -246,9 +247,9 @@ def compute_first_contact(self, dt: float, abs_tol: float = 1.0e-8) -> wp.array: outputs=[self._data._first_transition], device=self._device, ) - return self._data._first_transition + return self._data._first_transition_ta - def compute_first_air(self, dt: float, abs_tol: float = 1.0e-8) -> wp.array: + def compute_first_air(self, dt: float, abs_tol: float = 1.0e-8) -> ProxyArray: """Checks if sensors that have broken contact within the last :attr:`dt` seconds. This function checks if the sensors have broken contact within the last :attr:`dt` seconds @@ -288,7 +289,7 @@ def compute_first_air(self, dt: float, abs_tol: float = 1.0e-8) -> wp.array: outputs=[self._data._first_transition], device=self._device, ) - return self._data._first_transition + return self._data._first_transition_ta """ Implementation. diff --git a/source/isaaclab_newton/isaaclab_newton/sensors/contact_sensor/contact_sensor_data.py b/source/isaaclab_newton/isaaclab_newton/sensors/contact_sensor/contact_sensor_data.py index 7c8250ee808a..8ed00b91eab8 100644 --- a/source/isaaclab_newton/isaaclab_newton/sensors/contact_sensor/contact_sensor_data.py +++ b/source/isaaclab_newton/isaaclab_newton/sensors/contact_sensor/contact_sensor_data.py @@ -11,6 +11,7 @@ import warp as wp from isaaclab.sensors.contact_sensor.base_contact_sensor_data import BaseContactSensorData +from isaaclab.utils.warp import ProxyArray logger = logging.getLogger(__name__) @@ -32,12 +33,12 @@ class ContactSensorData(BaseContactSensorData): _first_transition: wp.array | None @property - def pose_w(self) -> wp.array | None: + def pose_w(self) -> ProxyArray | None: """Not supported by Newton contact sensor.""" raise NotImplementedError("pose_w is not supported by the Newton contact sensor.") @property - def pos_w(self) -> wp.array | None: + def pos_w(self) -> ProxyArray | None: """Position of the sensor origin in world frame. `wp.vec3f` array whose shape is (N,) where N is the number of sensors. Note, that when casted to as a @@ -46,10 +47,14 @@ def pos_w(self) -> wp.array | None: Note: If the :attr:`ContactSensorCfg.track_pose` is False, then this quantity is None. """ - return self._pos_w + if self._pos_w is None: + return None + if self._pos_w_ta is None: + self._pos_w_ta = ProxyArray(self._pos_w) + return self._pos_w_ta @property - def quat_w(self) -> wp.array | None: + def quat_w(self) -> ProxyArray | None: """Orientation of the sensor origin in quaternion (x, y, z, w) in world frame. `wp.quatf` whose shape is (N,) where N is the number of sensors. Note, that when casted to as a `torch.Tensor`, @@ -58,10 +63,14 @@ def quat_w(self) -> wp.array | None: Note: If the :attr:`ContactSensorCfg.track_pose` is False, then this quantity is None. """ - return self._quat_w + if self._quat_w is None: + return None + if self._quat_w_ta is None: + self._quat_w_ta = ProxyArray(self._quat_w) + return self._quat_w_ta @property - def net_forces_w(self) -> wp.array2d | None: + def net_forces_w(self) -> ProxyArray | None: """The net (total) contact forces in world frame. `wp.vec3f` array whose shape is (N, S) where N is the number of environments and S is the number of sensors. @@ -71,10 +80,12 @@ def net_forces_w(self) -> wp.array2d | None: This quantity is the sum of the contact forces acting on each sensor. It must not be confused with the total contact forces acting on the sensors (which also includes the tangential forces). """ - return self._net_forces_w + if self._net_forces_w is None: + return None + return self._net_forces_w_ta @property - def net_forces_w_history(self) -> wp.array3d | None: + def net_forces_w_history(self) -> ProxyArray | None: """The net (total) contact forces in world frame. `wp.vec3f` array whose shape is (N, T, S) where N is the number of environments, T is the configured history @@ -87,10 +98,12 @@ def net_forces_w_history(self) -> wp.array3d | None: This quantity is the sum of the contact forces acting on each sensor. It must not be confused with the total contact forces acting on the sensors (which also includes the tangential forces). """ - return self._net_forces_w_history + if self._net_forces_w_history is None: + return None + return self._net_forces_w_history_ta @property - def force_matrix_w(self) -> wp.array3d | None: + def force_matrix_w(self) -> ProxyArray | None: """The contact forces between sensors and filter objects in world frame. `wp.vec3f` array whose shape is (N, S, F) where N is the number of environments, S is number of sensors @@ -100,10 +113,14 @@ def force_matrix_w(self) -> wp.array3d | None: Note: If the :attr:`ContactSensorCfg.filter_prim_paths_expr` is empty, then this quantity is None. """ - return self._force_matrix_w + if self._force_matrix_w is None: + return None + if self._force_matrix_w_ta is None: + self._force_matrix_w_ta = ProxyArray(self._force_matrix_w) + return self._force_matrix_w_ta @property - def force_matrix_w_history(self) -> wp.array4d | None: + def force_matrix_w_history(self) -> ProxyArray | None: """The contact forces between sensors and filter objects in world frame. `wp.vec3f` array whose shape is (N, T, S, F) where N is the number of environments, T is the configured history @@ -115,20 +132,24 @@ def force_matrix_w_history(self) -> wp.array4d | None: Note: If the :attr:`ContactSensorCfg.filter_prim_paths_expr` is empty, then this quantity is None. """ - return self._force_matrix_w_history + if self._force_matrix_w_history is None: + return None + if self._force_matrix_w_history_ta is None: + self._force_matrix_w_history_ta = ProxyArray(self._force_matrix_w_history) + return self._force_matrix_w_history_ta @property - def contact_pos_w(self) -> wp.array | None: + def contact_pos_w(self) -> ProxyArray | None: """Not supported by Newton contact sensor.""" raise NotImplementedError("contact_pos_w is not supported by the Newton contact sensor.") @property - def friction_forces_w(self) -> wp.array | None: + def friction_forces_w(self) -> ProxyArray | None: """Not supported by Newton contact sensor.""" raise NotImplementedError("friction_forces_w is not supported by the Newton contact sensor.") @property - def last_air_time(self) -> wp.array2d | None: + def last_air_time(self) -> ProxyArray | None: """Time spent (in s) in the air before the last contact. `wp.float32` array whose shape is (N, S) where N is the number of environments and S is the number of sensors. @@ -137,10 +158,12 @@ def last_air_time(self) -> wp.array2d | None: Note: If the :attr:`ContactSensorCfg.track_air_time` is False, then this quantity is None. """ - return self._last_air_time + if self._last_air_time is None: + return None + return self._last_air_time_ta @property - def current_air_time(self) -> wp.array2d | None: + def current_air_time(self) -> ProxyArray | None: """Time spent (in s) in the air since the last detach. `wp.float32` array whose shape is (N, S) where N is the number of environments and S is the number of sensors. @@ -149,10 +172,12 @@ def current_air_time(self) -> wp.array2d | None: Note: If the :attr:`ContactSensorCfg.track_air_time` is False, then this quantity is None. """ - return self._current_air_time + if self._current_air_time is None: + return None + return self._current_air_time_ta @property - def last_contact_time(self) -> wp.array2d | None: + def last_contact_time(self) -> ProxyArray | None: """Time spent (in s) in contact before the last detach. `wp.float32` array whose shape is (N, S) where N is the number of environments and S is the number of sensors. @@ -161,10 +186,12 @@ def last_contact_time(self) -> wp.array2d | None: Note: If the :attr:`ContactSensorCfg.track_air_time` is False, then this quantity is None. """ - return self._last_contact_time + if self._last_contact_time is None: + return None + return self._last_contact_time_ta @property - def current_contact_time(self) -> wp.array2d | None: + def current_contact_time(self) -> ProxyArray | None: """Time spent (in s) in contact since the last contact. `wp.float32` array whose shape is (N, S) where N is the number of environments and S is the number of sensors. @@ -173,7 +200,9 @@ def current_contact_time(self) -> wp.array2d | None: Note: If the :attr:`ContactSensorCfg.track_air_time` is False, then this quantity is None. """ - return self._current_contact_time + if self._current_contact_time is None: + return None + return self._current_contact_time_ta def create_buffers( self, @@ -240,9 +269,31 @@ def create_buffers( self._last_contact_time = wp.zeros((num_envs, num_sensors), dtype=wp.float32, device=device) self._current_contact_time = wp.zeros((num_envs, num_sensors), dtype=wp.float32, device=device) self._first_transition = wp.zeros((num_envs, num_sensors), dtype=wp.float32, device=device) + self._first_transition_ta = ProxyArray(self._first_transition) else: self._last_air_time = None self._current_air_time = None self._last_contact_time = None self._current_contact_time = None self._first_transition = None + self._first_transition_ta = None + + # -- Pin ProxyArray instances for pre-allocated buffers + self._net_forces_w_ta = ProxyArray(self._net_forces_w) + self._net_forces_w_history_ta = ( + ProxyArray(self._net_forces_w_history) if self._net_forces_w_history is not None else None + ) + # -- Lazy ProxyArray instances for nullable buffers (pinned on first access) + self._pos_w_ta: ProxyArray | None = None + self._quat_w_ta: ProxyArray | None = None + self._force_matrix_w_ta: ProxyArray | None = None + self._force_matrix_w_history_ta: ProxyArray | None = None + # -- Pin ProxyArray instances for air/contact time buffers (eagerly when allocated) + self._last_air_time_ta = ProxyArray(self._last_air_time) if self._last_air_time is not None else None + self._current_air_time_ta = ProxyArray(self._current_air_time) if self._current_air_time is not None else None + self._last_contact_time_ta = ( + ProxyArray(self._last_contact_time) if self._last_contact_time is not None else None + ) + self._current_contact_time_ta = ( + ProxyArray(self._current_contact_time) if self._current_contact_time is not None else None + ) diff --git a/source/isaaclab_newton/isaaclab_newton/sensors/frame_transformer/frame_transformer_data.py b/source/isaaclab_newton/isaaclab_newton/sensors/frame_transformer/frame_transformer_data.py index 090a1642fbde..2b5f4f48006b 100644 --- a/source/isaaclab_newton/isaaclab_newton/sensors/frame_transformer/frame_transformer_data.py +++ b/source/isaaclab_newton/isaaclab_newton/sensors/frame_transformer/frame_transformer_data.py @@ -8,6 +8,7 @@ import warp as wp from isaaclab.sensors.frame_transformer.base_frame_transformer_data import BaseFrameTransformerData +from isaaclab.utils.warp import ProxyArray from isaaclab.utils.warp.math_ops import transform_to_vec_quat @@ -50,46 +51,46 @@ def target_frame_names(self) -> list[str]: return self._target_frame_names @property - def target_pose_source(self) -> wp.array: + def target_pose_source(self) -> ProxyArray: """Target poses relative to source frame, shape ``(num_envs, num_targets)``, dtype ``wp.transformf``.""" - return self._target_transforms + return ProxyArray(self._target_transforms) @property - def target_pos_source(self) -> wp.array: + def target_pos_source(self) -> ProxyArray: """Position of target frames relative to source frame [m], shape ``(num_envs, num_targets, 3)``.""" - return self._target_pos_source + return ProxyArray(self._target_pos_source) @property - def target_quat_source(self) -> wp.array: + def target_quat_source(self) -> ProxyArray: """Orientation of target frames relative to source frame (xyzw), shape ``(num_envs, num_targets, 4)``.""" - return self._target_quat_source + return ProxyArray(self._target_quat_source) @property - def target_pose_w(self) -> wp.array: + def target_pose_w(self) -> ProxyArray: """Target poses in world frame, shape ``(num_envs, num_targets)``, dtype ``wp.transformf``.""" - return self._target_transforms_w + return ProxyArray(self._target_transforms_w) @property - def target_pos_w(self) -> wp.array: + def target_pos_w(self) -> ProxyArray: """Position of target frames in world frame [m], shape ``(num_envs, num_targets, 3)``.""" - return self._target_pos_w + return ProxyArray(self._target_pos_w) @property - def target_quat_w(self) -> wp.array: + def target_quat_w(self) -> ProxyArray: """Orientation of target frames in world frame (xyzw), shape ``(num_envs, num_targets, 4)``.""" - return self._target_quat_w + return ProxyArray(self._target_quat_w) @property - def source_pose_w(self) -> wp.array: + def source_pose_w(self) -> ProxyArray: """Source pose in world frame, shape ``(num_envs,)``, dtype ``wp.transformf``.""" - return self._source_transforms + return ProxyArray(self._source_transforms) @property - def source_pos_w(self) -> wp.array: + def source_pos_w(self) -> ProxyArray: """Position of source frame in world frame [m], shape ``(num_envs, 3)``.""" - return self._source_pos_w + return ProxyArray(self._source_pos_w) @property - def source_quat_w(self) -> wp.array: + def source_quat_w(self) -> ProxyArray: """Orientation of source frame in world frame (xyzw), shape ``(num_envs, 4)``.""" - return self._source_quat_w + return ProxyArray(self._source_quat_w) diff --git a/source/isaaclab_newton/isaaclab_newton/sensors/imu/imu_data.py b/source/isaaclab_newton/isaaclab_newton/sensors/imu/imu_data.py index 25dc5dd66288..f288df131a2f 100644 --- a/source/isaaclab_newton/isaaclab_newton/sensors/imu/imu_data.py +++ b/source/isaaclab_newton/isaaclab_newton/sensors/imu/imu_data.py @@ -8,6 +8,7 @@ import warp as wp from isaaclab.sensors.imu import BaseImuData +from isaaclab.utils.warp import ProxyArray class ImuData(BaseImuData): @@ -16,19 +17,25 @@ class ImuData(BaseImuData): def __init__(self): self._ang_vel_b: wp.array | None = None self._lin_acc_b: wp.array | None = None + self._ang_vel_b_ta: ProxyArray | None = None + self._lin_acc_b_ta: ProxyArray | None = None @property - def ang_vel_b(self) -> wp.array | None: + def ang_vel_b(self) -> ProxyArray | None: """IMU frame angular velocity relative to the world expressed in IMU frame [rad/s]. Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). ``None`` before the simulation is initialized. """ - return self._ang_vel_b + if self._ang_vel_b is None: + return None + if self._ang_vel_b_ta is None: + self._ang_vel_b_ta = ProxyArray(self._ang_vel_b) + return self._ang_vel_b_ta @property - def lin_acc_b(self) -> wp.array | None: + def lin_acc_b(self) -> ProxyArray | None: """Linear acceleration (proper) in the IMU frame [m/s^2]. Zero in freefall, +g upward at rest. @@ -37,7 +44,11 @@ def lin_acc_b(self) -> wp.array | None: ``None`` before the simulation is initialized. """ - return self._lin_acc_b + if self._lin_acc_b is None: + return None + if self._lin_acc_b_ta is None: + self._lin_acc_b_ta = ProxyArray(self._lin_acc_b) + return self._lin_acc_b_ta def create_buffers(self, num_envs: int, device: str) -> None: """Create internal buffers for sensor data. @@ -48,3 +59,5 @@ def create_buffers(self, num_envs: int, device: str) -> None: """ self._ang_vel_b = wp.zeros(num_envs, dtype=wp.vec3f, device=device) self._lin_acc_b = wp.zeros(num_envs, dtype=wp.vec3f, device=device) + self._ang_vel_b_ta = None + self._lin_acc_b_ta = None diff --git a/source/isaaclab_newton/isaaclab_newton/sensors/pva/pva.py b/source/isaaclab_newton/isaaclab_newton/sensors/pva/pva.py index b46552ce4039..c000e17c437a 100644 --- a/source/isaaclab_newton/isaaclab_newton/sensors/pva/pva.py +++ b/source/isaaclab_newton/isaaclab_newton/sensors/pva/pva.py @@ -209,18 +209,16 @@ def _debug_vis_callback(self, event): if self._newton_model is None: return # base position (offset upward for visibility) - base_pos_w = wp.to_torch(self._data.pos_w).clone() + base_pos_w = self._data.pos_w.torch.clone() base_pos_w[:, 2] += 0.5 # arrow scale default_scale = self.acceleration_visualizer.cfg.markers["arrow"].scale - arrow_scale = torch.tensor(default_scale, device=self.device).repeat( - wp.to_torch(self._data.lin_acc_b).shape[0], 1 - ) + arrow_scale = torch.tensor(default_scale, device=self.device).repeat(self._data.lin_acc_b.torch.shape[0], 1) # arrow direction from acceleration up_axis = UsdGeom.GetStageUpAxis(self.stage) - pos_w_torch = wp.to_torch(self._data.pos_w) - quat_w_torch = wp.to_torch(self._data.quat_w) - lin_acc_b_torch = wp.to_torch(self._data.lin_acc_b) + pos_w_torch = self._data.pos_w.torch + quat_w_torch = self._data.quat_w.torch + lin_acc_b_torch = self._data.lin_acc_b.torch quat_opengl = math_utils.quat_from_matrix( math_utils.create_rotation_matrix_from_view( pos_w_torch, diff --git a/source/isaaclab_newton/isaaclab_newton/sensors/pva/pva_data.py b/source/isaaclab_newton/isaaclab_newton/sensors/pva/pva_data.py index ef80237a2e23..5b75a0087a26 100644 --- a/source/isaaclab_newton/isaaclab_newton/sensors/pva/pva_data.py +++ b/source/isaaclab_newton/isaaclab_newton/sensors/pva/pva_data.py @@ -8,6 +8,7 @@ import warp as wp from isaaclab.sensors.pva import BasePvaData +from isaaclab.utils.warp import ProxyArray class PvaData(BasePvaData): @@ -22,9 +23,18 @@ def __init__(self): self._ang_vel_b: wp.array | None = None self._lin_acc_b: wp.array | None = None self._ang_acc_b: wp.array | None = None + # ProxyArray caches + self._pose_w_ta: ProxyArray | None = None + self._pos_w_ta: ProxyArray | None = None + self._quat_w_ta: ProxyArray | None = None + self._projected_gravity_b_ta: ProxyArray | None = None + self._lin_vel_b_ta: ProxyArray | None = None + self._ang_vel_b_ta: ProxyArray | None = None + self._lin_acc_b_ta: ProxyArray | None = None + self._ang_acc_b_ta: ProxyArray | None = None @property - def pose_w(self) -> wp.array | None: + def pose_w(self) -> ProxyArray | None: """Pose of the sensor origin in world frame [m, unitless]. Shape is (num_instances,), dtype = wp.transformf. In torch this resolves to (num_instances, 7). @@ -32,20 +42,28 @@ def pose_w(self) -> wp.array | None: ``None`` before the simulation is initialized. """ - return self._pose_w + if self._pose_w is None: + return None + if self._pose_w_ta is None: + self._pose_w_ta = ProxyArray(self._pose_w) + return self._pose_w_ta @property - def pos_w(self) -> wp.array | None: + def pos_w(self) -> ProxyArray | None: """Position of the sensor origin in world frame [m]. Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). ``None`` before the simulation is initialized. """ - return self._pos_w + if self._pos_w is None: + return None + if self._pos_w_ta is None: + self._pos_w_ta = ProxyArray(self._pos_w) + return self._pos_w_ta @property - def quat_w(self) -> wp.array | None: + def quat_w(self) -> ProxyArray | None: """Orientation of the sensor origin in world frame. Shape is (num_instances,), dtype = wp.quatf. In torch this resolves to (num_instances, 4). @@ -53,40 +71,56 @@ def quat_w(self) -> wp.array | None: ``None`` before the simulation is initialized. """ - return self._quat_w + if self._quat_w is None: + return None + if self._quat_w_ta is None: + self._quat_w_ta = ProxyArray(self._quat_w) + return self._quat_w_ta @property - def projected_gravity_b(self) -> wp.array | None: + def projected_gravity_b(self) -> ProxyArray | None: """Gravity direction unit vector projected on the PVA frame [unitless]. Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). ``None`` before the simulation is initialized. """ - return self._projected_gravity_b + if self._projected_gravity_b is None: + return None + if self._projected_gravity_b_ta is None: + self._projected_gravity_b_ta = ProxyArray(self._projected_gravity_b) + return self._projected_gravity_b_ta @property - def lin_vel_b(self) -> wp.array | None: + def lin_vel_b(self) -> ProxyArray | None: """PVA frame linear velocity relative to the world expressed in PVA frame [m/s]. Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). ``None`` before the simulation is initialized. """ - return self._lin_vel_b + if self._lin_vel_b is None: + return None + if self._lin_vel_b_ta is None: + self._lin_vel_b_ta = ProxyArray(self._lin_vel_b) + return self._lin_vel_b_ta @property - def ang_vel_b(self) -> wp.array | None: + def ang_vel_b(self) -> ProxyArray | None: """PVA frame angular velocity relative to the world expressed in PVA frame [rad/s]. Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). ``None`` before the simulation is initialized. """ - return self._ang_vel_b + if self._ang_vel_b is None: + return None + if self._ang_vel_b_ta is None: + self._ang_vel_b_ta = ProxyArray(self._ang_vel_b) + return self._ang_vel_b_ta @property - def lin_acc_b(self) -> wp.array | None: + def lin_acc_b(self) -> ProxyArray | None: """Linear acceleration (coordinate) in the PVA frame [m/s^2]. Equal to -g in freefall, zero at rest. @@ -95,17 +129,25 @@ def lin_acc_b(self) -> wp.array | None: ``None`` before the simulation is initialized. """ - return self._lin_acc_b + if self._lin_acc_b is None: + return None + if self._lin_acc_b_ta is None: + self._lin_acc_b_ta = ProxyArray(self._lin_acc_b) + return self._lin_acc_b_ta @property - def ang_acc_b(self) -> wp.array | None: + def ang_acc_b(self) -> ProxyArray | None: """PVA frame angular acceleration relative to the world expressed in PVA frame [rad/s^2]. Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). ``None`` before the simulation is initialized. """ - return self._ang_acc_b + if self._ang_acc_b is None: + return None + if self._ang_acc_b_ta is None: + self._ang_acc_b_ta = ProxyArray(self._ang_acc_b) + return self._ang_acc_b_ta def create_buffers(self, num_envs: int, device: str) -> None: """Create internal buffers for sensor data. @@ -122,3 +164,12 @@ def create_buffers(self, num_envs: int, device: str) -> None: self._ang_vel_b = wp.zeros(num_envs, dtype=wp.vec3f, device=device) self._lin_acc_b = wp.zeros(num_envs, dtype=wp.vec3f, device=device) self._ang_acc_b = wp.zeros(num_envs, dtype=wp.vec3f, device=device) + # Reset ProxyArray caches + self._pose_w_ta = None + self._pos_w_ta = None + self._quat_w_ta = None + self._projected_gravity_b_ta = None + self._lin_vel_b_ta = None + self._ang_vel_b_ta = None + self._lin_acc_b_ta = None + self._ang_acc_b_ta = None diff --git a/source/isaaclab_newton/isaaclab_newton/sim/views/newton_site_frame_view.py b/source/isaaclab_newton/isaaclab_newton/sim/views/newton_site_frame_view.py index e4f2285cb528..c7758f0cbe42 100644 --- a/source/isaaclab_newton/isaaclab_newton/sim/views/newton_site_frame_view.py +++ b/source/isaaclab_newton/isaaclab_newton/sim/views/newton_site_frame_view.py @@ -16,6 +16,7 @@ import isaaclab.sim as sim_utils from isaaclab.physics import PhysicsEvent from isaaclab.sim.views.base_frame_view import BaseFrameView +from isaaclab.utils.warp import ProxyArray from isaaclab_newton.physics.newton_manager import NewtonManager @@ -506,7 +507,7 @@ class NewtonSiteFrameView(BaseFrameView): ``set_world_poses`` and ``set_local_poses`` update ``site_local`` -- neither touches ``body_q``. - All getters return ``wp.array``. Setters accept ``wp.array``. + Pose getters return :class:`~isaaclab.utils.warp.ProxyArray`. Setters accept ``wp.array``. Raises: ValueError: If any matched prim resolves to a Newton physics body @@ -622,6 +623,10 @@ def _initialize_impl(self, model) -> None: self._quat_buf = wp.zeros(self.count, dtype=wp.vec4f, device=device) self._local_pos_buf = wp.zeros(self.count, dtype=wp.vec3f, device=device) self._local_quat_buf = wp.zeros(self.count, dtype=wp.vec4f, device=device) + self._pos_ta = ProxyArray(self._pos_buf) + self._quat_ta = ProxyArray(self._quat_buf) + self._local_pos_ta = ProxyArray(self._local_pos_buf) + self._local_quat_ta = ProxyArray(self._local_quat_buf) @staticmethod def _resolve_ancestor_body( @@ -668,19 +673,25 @@ def count(self) -> int: """Number of prims in this view.""" return len(self._prims) + @property + def device(self) -> str: + """Device where arrays are allocated (cpu or cuda).""" + return self._device + # ------------------------------------------------------------------ # World poses # ------------------------------------------------------------------ - def get_world_poses(self, indices: wp.array | None = None) -> tuple[wp.array, wp.array]: + def get_world_poses(self, indices: wp.array | None = None) -> tuple[ProxyArray, ProxyArray]: """Get world-space positions and orientations. Args: indices: Subset of sites to query. ``None`` means all sites. Returns: - A tuple ``(positions, orientations)`` as ``wp.array`` of shapes - ``(M, 3)`` and ``(M, 4)`` respectively. + A tuple ``(positions, orientations)`` of :class:`~isaaclab.utils.warp.ProxyArray` + wrappers. Use ``.warp`` for the underlying ``wp.array`` or ``.torch`` for a + cached zero-copy ``torch.Tensor`` view. """ state = NewtonManager.get_state_0() @@ -695,7 +706,7 @@ def get_world_poses(self, indices: wp.array | None = None) -> tuple[wp.array, wp outputs=[pos_buf, quat_buf], device=self._device, ) - return pos_buf, quat_buf + return ProxyArray(pos_buf), ProxyArray(quat_buf) wp.launch( _compute_site_world_transforms, @@ -704,7 +715,7 @@ def get_world_poses(self, indices: wp.array | None = None) -> tuple[wp.array, wp outputs=[self._pos_buf, self._quat_buf], device=self._device, ) - return self._pos_buf, self._quat_buf + return self._pos_ta, self._quat_ta def set_world_poses( self, @@ -731,11 +742,11 @@ def set_world_poses( state = NewtonManager.get_state_0() if positions is None or orientations is None: - cur_pos, cur_quat = self.get_world_poses(indices) + cur_pos_ta, cur_quat_ta = self.get_world_poses(indices) if positions is None: - positions = cur_pos + positions = cur_pos_ta.warp if orientations is None: - orientations = cur_quat + orientations = cur_quat_ta.warp if indices is not None: wp.launch( @@ -756,7 +767,7 @@ def set_world_poses( # Local poses (parent-relative) # ------------------------------------------------------------------ - def get_local_poses(self, indices: wp.array | None = None) -> tuple[wp.array, wp.array]: + def get_local_poses(self, indices: wp.array | None = None) -> tuple[ProxyArray, ProxyArray]: """Get parent-relative positions and orientations. Computes ``inv(parent_world) * prim_world`` for each site. @@ -765,8 +776,9 @@ def get_local_poses(self, indices: wp.array | None = None) -> tuple[wp.array, wp indices: Subset of sites to query. ``None`` means all sites. Returns: - A tuple ``(translations, orientations)`` as ``wp.array`` of shapes - ``(M, 3)`` and ``(M, 4)`` respectively. + A tuple ``(translations, orientations)`` of :class:`~isaaclab.utils.warp.ProxyArray` + wrappers. Use ``.warp`` for the underlying ``wp.array`` or ``.torch`` for a + cached zero-copy ``torch.Tensor`` view. """ state = NewtonManager.get_state_0() @@ -788,7 +800,7 @@ def get_local_poses(self, indices: wp.array | None = None) -> tuple[wp.array, wp outputs=[pos_buf, quat_buf], device=self._device, ) - return pos_buf, quat_buf + return ProxyArray(pos_buf), ProxyArray(quat_buf) wp.launch( _compute_site_local_transforms, @@ -803,7 +815,7 @@ def get_local_poses(self, indices: wp.array | None = None) -> tuple[wp.array, wp outputs=[self._local_pos_buf, self._local_quat_buf], device=self._device, ) - return self._local_pos_buf, self._local_quat_buf + return self._local_pos_ta, self._local_quat_ta def set_local_poses( self, @@ -830,11 +842,11 @@ def set_local_poses( state = NewtonManager.get_state_0() if translations is None or orientations is None: - cur_pos, cur_quat = self.get_local_poses(indices) + cur_pos_ta, cur_quat_ta = self.get_local_poses(indices) if translations is None: - translations = cur_pos + translations = cur_pos_ta.warp if orientations is None: - orientations = cur_quat + orientations = cur_quat_ta.warp if indices is not None: wp.launch( diff --git a/source/isaaclab_newton/test/assets/test_articulation.py b/source/isaaclab_newton/test/assets/test_articulation.py index a03a596dd690..5a4e77fd9eaf 100644 --- a/source/isaaclab_newton/test/assets/test_articulation.py +++ b/source/isaaclab_newton/test/assets/test_articulation.py @@ -412,9 +412,9 @@ def test_initialization_floating_base_non_root(sim, num_articulations, device, a # Check that is fixed base assert not articulation.is_fixed_base # Check buffers that exists and have correct shapes - assert wp.to_torch(articulation.data.root_pos_w).shape == (num_articulations, 3) - assert wp.to_torch(articulation.data.root_quat_w).shape == (num_articulations, 4) - assert wp.to_torch(articulation.data.joint_pos).shape == (num_articulations, 21) + assert articulation.data.root_pos_w.torch.shape == (num_articulations, 3) + assert articulation.data.root_quat_w.torch.shape == (num_articulations, 4) + assert articulation.data.joint_pos.torch.shape == (num_articulations, 21) # -- actuator type for actuator_name, actuator in articulation.actuators.items(): @@ -461,11 +461,11 @@ def test_initialization_floating_base(sim, num_articulations, device, add_ground # Check that floating base assert not articulation.is_fixed_base # Check buffers that exists and have correct shapes - assert wp.to_torch(articulation.data.root_pos_w).shape == (num_articulations, 3) - assert wp.to_torch(articulation.data.root_quat_w).shape == (num_articulations, 4) - assert wp.to_torch(articulation.data.joint_pos).shape == (num_articulations, 12) - assert wp.to_torch(articulation.data.body_mass).shape == (num_articulations, articulation.num_bodies) - assert wp.to_torch(articulation.data.body_inertia).shape == (num_articulations, articulation.num_bodies, 9) + assert articulation.data.root_pos_w.torch.shape == (num_articulations, 3) + assert articulation.data.root_quat_w.torch.shape == (num_articulations, 4) + assert articulation.data.joint_pos.torch.shape == (num_articulations, 12) + assert articulation.data.body_mass.torch.shape == (num_articulations, articulation.num_bodies) + assert articulation.data.body_inertia.torch.shape == (num_articulations, articulation.num_bodies, 9) # -- actuator type for actuator_name, actuator in articulation.actuators.items(): @@ -511,11 +511,11 @@ def test_initialization_fixed_base(sim, num_articulations, device, articulation_ # Check that fixed base assert articulation.is_fixed_base # Check buffers that exists and have correct shapes - assert wp.to_torch(articulation.data.root_pos_w).shape == (num_articulations, 3) - assert wp.to_torch(articulation.data.root_quat_w).shape == (num_articulations, 4) - assert wp.to_torch(articulation.data.joint_pos).shape == (num_articulations, 9) - assert wp.to_torch(articulation.data.body_mass).shape == (num_articulations, articulation.num_bodies) - assert wp.to_torch(articulation.data.body_inertia).shape == (num_articulations, articulation.num_bodies, 9) + assert articulation.data.root_pos_w.torch.shape == (num_articulations, 3) + assert articulation.data.root_quat_w.torch.shape == (num_articulations, 4) + assert articulation.data.joint_pos.torch.shape == (num_articulations, 9) + assert articulation.data.body_mass.torch.shape == (num_articulations, articulation.num_bodies) + assert articulation.data.body_inertia.torch.shape == (num_articulations, articulation.num_bodies, 9) # -- actuator type for actuator_name, actuator in articulation.actuators.items(): @@ -530,12 +530,12 @@ def test_initialization_fixed_base(sim, num_articulations, device, articulation_ articulation.update(sim.cfg.dt) # check that the root is at the correct state - its default state as it is fixed base - default_root_pose = wp.to_torch(articulation.data.default_root_pose).clone() - default_root_vel = wp.to_torch(articulation.data.default_root_vel).clone() + default_root_pose = articulation.data.default_root_pose.torch.clone() + default_root_vel = articulation.data.default_root_vel.torch.clone() default_root_pose[:, :3] = default_root_pose[:, :3] + translations - torch.testing.assert_close(wp.to_torch(articulation.data.root_link_pose_w), default_root_pose) - torch.testing.assert_close(wp.to_torch(articulation.data.root_com_vel_w), default_root_vel) + torch.testing.assert_close(articulation.data.root_link_pose_w.torch, default_root_pose) + torch.testing.assert_close(articulation.data.root_com_vel_w.torch, default_root_vel) @pytest.mark.isaacsim_ci @@ -570,11 +570,11 @@ def test_initialization_fixed_base_single_joint(sim, num_articulations, device, # Check that fixed base assert articulation.is_fixed_base # Check buffers that exists and have correct shapes - assert wp.to_torch(articulation.data.root_pos_w).shape == (num_articulations, 3) - assert wp.to_torch(articulation.data.root_quat_w).shape == (num_articulations, 4) - assert wp.to_torch(articulation.data.joint_pos).shape == (num_articulations, 1) - assert wp.to_torch(articulation.data.body_mass).shape == (num_articulations, articulation.num_bodies) - assert wp.to_torch(articulation.data.body_inertia).shape == (num_articulations, articulation.num_bodies, 9) + assert articulation.data.root_pos_w.torch.shape == (num_articulations, 3) + assert articulation.data.root_quat_w.torch.shape == (num_articulations, 4) + assert articulation.data.joint_pos.torch.shape == (num_articulations, 1) + assert articulation.data.body_mass.torch.shape == (num_articulations, articulation.num_bodies) + assert articulation.data.body_inertia.torch.shape == (num_articulations, articulation.num_bodies, 9) # -- actuator type for actuator_name, actuator in articulation.actuators.items(): @@ -589,12 +589,12 @@ def test_initialization_fixed_base_single_joint(sim, num_articulations, device, articulation.update(sim.cfg.dt) # check that the root is at the correct state - its default state as it is fixed base - default_root_pose = wp.to_torch(articulation.data.default_root_pose).clone() - default_root_vel = wp.to_torch(articulation.data.default_root_vel).clone() + default_root_pose = articulation.data.default_root_pose.torch.clone() + default_root_vel = articulation.data.default_root_vel.torch.clone() default_root_pose[:, :3] = default_root_pose[:, :3] + translations - torch.testing.assert_close(wp.to_torch(articulation.data.root_link_pose_w), default_root_pose) - torch.testing.assert_close(wp.to_torch(articulation.data.root_com_vel_w), default_root_vel) + torch.testing.assert_close(articulation.data.root_link_pose_w.torch, default_root_pose) + torch.testing.assert_close(articulation.data.root_com_vel_w.torch, default_root_vel) @pytest.mark.isaacsim_ci @@ -628,11 +628,11 @@ def test_initialization_hand_with_tendons(sim, num_articulations, device, articu # Check that fixed base assert articulation.is_fixed_base # Check buffers that exists and have correct shapes - assert wp.to_torch(articulation.data.root_pos_w).shape == (num_articulations, 3) - assert wp.to_torch(articulation.data.root_quat_w).shape == (num_articulations, 4) - assert wp.to_torch(articulation.data.joint_pos).shape == (num_articulations, 24) - assert wp.to_torch(articulation.data.body_mass).shape == (num_articulations, articulation.num_bodies) - assert wp.to_torch(articulation.data.body_inertia).shape == (num_articulations, articulation.num_bodies, 9) + assert articulation.data.root_pos_w.torch.shape == (num_articulations, 3) + assert articulation.data.root_quat_w.torch.shape == (num_articulations, 4) + assert articulation.data.joint_pos.torch.shape == (num_articulations, 24) + assert articulation.data.body_mass.torch.shape == (num_articulations, articulation.num_bodies) + assert articulation.data.body_inertia.torch.shape == (num_articulations, articulation.num_bodies, 9) # -- actuator type for actuator_name, actuator in articulation.actuators.items(): @@ -682,9 +682,9 @@ def test_initialization_floating_base_made_fixed_base( # Check that is fixed base assert articulation.is_fixed_base # Check buffers that exists and have correct shapes - assert wp.to_torch(articulation.data.root_pos_w).shape == (num_articulations, 3) - assert wp.to_torch(articulation.data.root_quat_w).shape == (num_articulations, 4) - assert wp.to_torch(articulation.data.joint_pos).shape == (num_articulations, 12) + assert articulation.data.root_pos_w.torch.shape == (num_articulations, 3) + assert articulation.data.root_quat_w.torch.shape == (num_articulations, 4) + assert articulation.data.joint_pos.torch.shape == (num_articulations, 12) # Simulate physics for _ in range(10): @@ -694,12 +694,12 @@ def test_initialization_floating_base_made_fixed_base( articulation.update(sim.cfg.dt) # check that the root is at the correct state - its default state as it is fixed base - default_root_pose = wp.to_torch(articulation.data.default_root_pose).clone() - default_root_vel = wp.to_torch(articulation.data.default_root_vel).clone() + default_root_pose = articulation.data.default_root_pose.torch.clone() + default_root_vel = articulation.data.default_root_vel.torch.clone() default_root_pose[:, :3] = default_root_pose[:, :3] + translations - torch.testing.assert_close(wp.to_torch(articulation.data.root_link_pose_w), default_root_pose) - torch.testing.assert_close(wp.to_torch(articulation.data.root_com_vel_w), default_root_vel) + torch.testing.assert_close(articulation.data.root_link_pose_w.torch, default_root_pose) + torch.testing.assert_close(articulation.data.root_com_vel_w.torch, default_root_vel) @pytest.mark.isaacsim_ci @@ -737,9 +737,9 @@ def test_initialization_fixed_base_made_floating_base( # Check that is floating base assert not articulation.is_fixed_base # Check buffers that exists and have correct shapes - assert wp.to_torch(articulation.data.root_pos_w).shape == (num_articulations, 3) - assert wp.to_torch(articulation.data.root_quat_w).shape == (num_articulations, 4) - assert wp.to_torch(articulation.data.joint_pos).shape == (num_articulations, 9) + assert articulation.data.root_pos_w.torch.shape == (num_articulations, 3) + assert articulation.data.root_quat_w.torch.shape == (num_articulations, 4) + assert articulation.data.joint_pos.torch.shape == (num_articulations, 9) # Simulate physics for _ in range(10): @@ -835,7 +835,7 @@ def test_joint_pos_limits(sim, num_articulations, device, add_ground_plane, arti assert articulation.is_initialized # Get current default joint pos - default_joint_pos = wp.to_torch(articulation._data.default_joint_pos).clone() + default_joint_pos = articulation._data.default_joint_pos.torch.clone() # Set new joint limits limits = torch.zeros(num_articulations, articulation.num_joints, 2, device=device) @@ -844,8 +844,8 @@ def test_joint_pos_limits(sim, num_articulations, device, add_ground_plane, arti articulation.write_joint_position_limit_to_sim_index(limits=limits) # Check new limits are in place - torch.testing.assert_close(wp.to_torch(articulation._data.joint_pos_limits), limits) - torch.testing.assert_close(wp.to_torch(articulation._data.default_joint_pos), default_joint_pos) + torch.testing.assert_close(articulation._data.joint_pos_limits.torch, limits) + torch.testing.assert_close(articulation._data.default_joint_pos.torch, default_joint_pos) # Set new joint limits with indexing env_ids = torch.arange(1, device=device, dtype=torch.int32) @@ -856,8 +856,8 @@ def test_joint_pos_limits(sim, num_articulations, device, add_ground_plane, arti articulation.write_joint_position_limit_to_sim_index(limits=limits, env_ids=env_ids, joint_ids=joint_ids) # Check new limits are in place - torch.testing.assert_close(wp.to_torch(articulation._data.joint_pos_limits)[env_ids][:, joint_ids], limits) - torch.testing.assert_close(wp.to_torch(articulation._data.default_joint_pos), default_joint_pos) + torch.testing.assert_close(articulation._data.joint_pos_limits.torch[env_ids][:, joint_ids], limits) + torch.testing.assert_close(articulation._data.default_joint_pos.torch, default_joint_pos) # Set new joint limits that invalidate default joint pos limits = torch.zeros(num_articulations, articulation.num_joints, 2, device=device) @@ -866,7 +866,7 @@ def test_joint_pos_limits(sim, num_articulations, device, add_ground_plane, arti articulation.write_joint_position_limit_to_sim_index(limits=limits) # Check if all values are within the bounds - default_joint_pos_torch = wp.to_torch(articulation._data.default_joint_pos) + default_joint_pos_torch = articulation._data.default_joint_pos.torch within_bounds = (default_joint_pos_torch >= limits[..., 0]) & (default_joint_pos_torch <= limits[..., 1]) assert torch.all(within_bounds) @@ -877,7 +877,7 @@ def test_joint_pos_limits(sim, num_articulations, device, add_ground_plane, arti articulation.write_joint_position_limit_to_sim_index(limits=limits, env_ids=env_ids, joint_ids=joint_ids) # Check if all values are within the bounds - default_joint_pos_torch = wp.to_torch(articulation._data.default_joint_pos) + default_joint_pos_torch = articulation._data.default_joint_pos.torch within_bounds = (default_joint_pos_torch[env_ids][:, joint_ids] >= limits[..., 0]) & ( default_joint_pos_torch[env_ids][:, joint_ids] <= limits[..., 1] ) @@ -906,14 +906,14 @@ def __init__(self, art): assert articulation.is_initialized # Case A: no clipping → should NOT terminate - wp.to_torch(articulation._data.computed_torque).zero_() - wp.to_torch(articulation._data.applied_torque).zero_() + articulation._data.computed_torque.torch.zero_() + articulation._data.applied_torque.torch.zero_() out = joint_effort_out_of_limit(env, robot_all) # [N] assert torch.all(~out) # Case B: simulate clipping → should terminate - wp.to_torch(articulation._data.computed_torque).fill_(100.0) # pretend controller commanded 100 - wp.to_torch(articulation._data.applied_torque).fill_(50.0) # pretend actuator clipped to 50 + articulation._data.computed_torque.torch.fill_(100.0) # pretend controller commanded 100 + articulation._data.applied_torque.torch.fill_(50.0) # pretend actuator clipped to 50 out = joint_effort_out_of_limit(env, robot_all) # [N] assert torch.all(out) @@ -944,13 +944,13 @@ def test_external_force_buffer(sim, num_articulations, device, articulation_type body_ids, _ = articulation.find_bodies("base") # reset root state - articulation.write_root_pose_to_sim_index(root_pose=wp.to_torch(articulation.data.default_root_pose).clone()) - articulation.write_root_velocity_to_sim_index(root_velocity=wp.to_torch(articulation.data.default_root_vel).clone()) + articulation.write_root_pose_to_sim_index(root_pose=articulation.data.default_root_pose.torch.clone()) + articulation.write_root_velocity_to_sim_index(root_velocity=articulation.data.default_root_vel.torch.clone()) # reset dof state joint_pos, joint_vel = ( - wp.to_torch(articulation.data.default_joint_pos), - wp.to_torch(articulation.data.default_joint_vel), + articulation.data.default_joint_pos.torch, + articulation.data.default_joint_vel.torch, ) articulation.write_joint_position_to_sim_index(position=joint_pos) articulation.write_joint_velocity_to_sim_index(velocity=joint_vel) @@ -983,8 +983,8 @@ def test_external_force_buffer(sim, num_articulations, device, articulation_type # check if the articulation's force and torque buffers are correctly updated for i in range(num_articulations): - assert wp.to_torch(articulation.permanent_wrench_composer.composed_force)[i, 0, 0].item() == force - assert wp.to_torch(articulation.permanent_wrench_composer.composed_torque)[i, 0, 0].item() == force + assert articulation.permanent_wrench_composer.composed_force.torch[i, 0, 0].item() == force + assert articulation.permanent_wrench_composer.composed_torque.torch[i, 0, 0].item() == force # Check if the instantaneous wrench is correctly added to the permanent wrench articulation.instantaneous_wrench_composer.add_forces_and_torques_index( @@ -994,7 +994,7 @@ def test_external_force_buffer(sim, num_articulations, device, articulation_type ) # apply action to the articulation - articulation.set_joint_position_target_index(target=wp.to_torch(articulation.data.default_joint_pos).clone()) + articulation.set_joint_position_target_index(target=articulation.data.default_joint_pos.torch.clone()) articulation.write_data_to_sim() # perform step @@ -1034,14 +1034,12 @@ def test_external_force_on_single_body(sim, num_articulations, device, articulat # Now we are ready! for _ in range(5): # reset root state - articulation.write_root_pose_to_sim_index(root_pose=wp.to_torch(articulation.data.default_root_pose).clone()) - articulation.write_root_velocity_to_sim_index( - root_velocity=wp.to_torch(articulation.data.default_root_vel).clone() - ) + articulation.write_root_pose_to_sim_index(root_pose=articulation.data.default_root_pose.torch.clone()) + articulation.write_root_velocity_to_sim_index(root_velocity=articulation.data.default_root_vel.torch.clone()) # reset dof state joint_pos, joint_vel = ( - wp.to_torch(articulation.data.default_joint_pos), - wp.to_torch(articulation.data.default_joint_vel), + articulation.data.default_joint_pos.torch, + articulation.data.default_joint_vel.torch, ) articulation.write_joint_position_to_sim_index(position=joint_pos) articulation.write_joint_velocity_to_sim_index(velocity=joint_vel) @@ -1054,9 +1052,7 @@ def test_external_force_on_single_body(sim, num_articulations, device, articulat # perform simulation for _ in range(100): # apply action to the articulation - articulation.set_joint_position_target_index( - target=wp.to_torch(articulation.data.default_joint_pos).clone() - ) + articulation.set_joint_position_target_index(target=articulation.data.default_joint_pos.torch.clone()) articulation.write_data_to_sim() # perform step sim.step() @@ -1064,7 +1060,7 @@ def test_external_force_on_single_body(sim, num_articulations, device, articulat articulation.update(sim.cfg.dt) # check condition that the articulations have fallen down for i in range(num_articulations): - assert wp.to_torch(articulation.data.root_pos_w)[i, 2].item() < 0.2 + assert articulation.data.root_pos_w.torch[i, 2].item() < 0.2 @pytest.mark.isaacsim_ci @@ -1106,17 +1102,15 @@ def test_external_force_on_single_body_at_position(sim, num_articulations, devic # Now we are ready! for i in range(5): # reset root state - root_pose = wp.to_torch(articulation.data.default_root_pose).clone() + root_pose = articulation.data.default_root_pose.torch.clone() root_pose[0, 0] = 2.5 # space them apart by 2.5m articulation.write_root_pose_to_sim_index(root_pose=root_pose) - articulation.write_root_velocity_to_sim_index( - root_velocity=wp.to_torch(articulation.data.default_root_vel).clone() - ) + articulation.write_root_velocity_to_sim_index(root_velocity=articulation.data.default_root_vel.torch.clone()) # reset dof state joint_pos, joint_vel = ( - wp.to_torch(articulation.data.default_joint_pos), - wp.to_torch(articulation.data.default_joint_vel), + articulation.data.default_joint_pos.torch, + articulation.data.default_joint_vel.torch, ) articulation.write_joint_position_to_sim_index(position=joint_pos) articulation.write_joint_velocity_to_sim_index(velocity=joint_vel) @@ -1126,7 +1120,7 @@ def test_external_force_on_single_body_at_position(sim, num_articulations, devic is_global = False if i % 2 == 0: - body_com_pos_w = wp.to_torch(articulation.data.body_com_pos_w)[:, body_ids, :3] + body_com_pos_w = articulation.data.body_com_pos_w.torch[:, body_ids, :3] # is_global = True external_wrench_positions_b[..., 0] = 0.0 external_wrench_positions_b[..., 1] = 1.0 @@ -1154,9 +1148,7 @@ def test_external_force_on_single_body_at_position(sim, num_articulations, devic # perform simulation for _ in range(100): # apply action to the articulation - articulation.set_joint_position_target_index( - target=wp.to_torch(articulation.data.default_joint_pos).clone() - ) + articulation.set_joint_position_target_index(target=articulation.data.default_joint_pos.torch.clone()) articulation.write_data_to_sim() # perform step sim.step() @@ -1164,7 +1156,7 @@ def test_external_force_on_single_body_at_position(sim, num_articulations, devic articulation.update(sim.cfg.dt) # check condition that the articulations have fallen down for i in range(num_articulations): - assert wp.to_torch(articulation.data.root_pos_w)[i, 2].item() < 0.2 + assert articulation.data.root_pos_w.torch[i, 2].item() < 0.2 @pytest.mark.isaacsim_ci @@ -1198,14 +1190,12 @@ def test_external_force_on_multiple_bodies(sim, num_articulations, device, artic # Now we are ready! for _ in range(5): # reset root state - articulation.write_root_pose_to_sim_index(root_pose=wp.to_torch(articulation.data.default_root_pose).clone()) - articulation.write_root_velocity_to_sim_index( - root_velocity=wp.to_torch(articulation.data.default_root_vel).clone() - ) + articulation.write_root_pose_to_sim_index(root_pose=articulation.data.default_root_pose.torch.clone()) + articulation.write_root_velocity_to_sim_index(root_velocity=articulation.data.default_root_vel.torch.clone()) # reset dof state joint_pos, joint_vel = ( - wp.to_torch(articulation.data.default_joint_pos), - wp.to_torch(articulation.data.default_joint_vel), + articulation.data.default_joint_pos.torch, + articulation.data.default_joint_vel.torch, ) articulation.write_joint_position_to_sim_index(position=joint_pos) articulation.write_joint_velocity_to_sim_index(velocity=joint_vel) @@ -1218,9 +1208,7 @@ def test_external_force_on_multiple_bodies(sim, num_articulations, device, artic # perform simulation for _ in range(100): # apply action to the articulation - articulation.set_joint_position_target_index( - target=wp.to_torch(articulation.data.default_joint_pos).clone() - ) + articulation.set_joint_position_target_index(target=articulation.data.default_joint_pos.torch.clone()) articulation.write_data_to_sim() # perform step sim.step() @@ -1229,7 +1217,7 @@ def test_external_force_on_multiple_bodies(sim, num_articulations, device, artic # check condition for i in range(num_articulations): # since there is a moment applied on the articulation, the articulation should rotate - assert wp.to_torch(articulation.data.root_ang_vel_w)[i, 2].item() > 0.1 + assert articulation.data.root_ang_vel_w.torch[i, 2].item() > 0.1 @pytest.mark.isaacsim_ci @@ -1272,14 +1260,12 @@ def test_external_force_on_multiple_bodies_at_position(sim, num_articulations, d # Now we are ready! for i in range(5): # reset root state - articulation.write_root_pose_to_sim_index(root_pose=wp.to_torch(articulation.data.default_root_pose).clone()) - articulation.write_root_velocity_to_sim_index( - root_velocity=wp.to_torch(articulation.data.default_root_vel).clone() - ) + articulation.write_root_pose_to_sim_index(root_pose=articulation.data.default_root_pose.torch.clone()) + articulation.write_root_velocity_to_sim_index(root_velocity=articulation.data.default_root_vel.torch.clone()) # reset dof state joint_pos, joint_vel = ( - wp.to_torch(articulation.data.default_joint_pos), - wp.to_torch(articulation.data.default_joint_vel), + articulation.data.default_joint_pos.torch, + articulation.data.default_joint_vel.torch, ) articulation.write_joint_position_to_sim_index(position=joint_pos) articulation.write_joint_velocity_to_sim_index(velocity=joint_vel) @@ -1288,7 +1274,7 @@ def test_external_force_on_multiple_bodies_at_position(sim, num_articulations, d is_global = False if i % 2 == 0: - body_com_pos_w = wp.to_torch(articulation.data.body_com_pos_w)[:, body_ids, :3] + body_com_pos_w = articulation.data.body_com_pos_w.torch[:, body_ids, :3] is_global = True external_wrench_positions_b[..., 0] = 0.0 external_wrench_positions_b[..., 1] = 1.0 @@ -1317,9 +1303,7 @@ def test_external_force_on_multiple_bodies_at_position(sim, num_articulations, d # perform simulation for _ in range(100): # apply action to the articulation - articulation.set_joint_position_target_index( - target=wp.to_torch(articulation.data.default_joint_pos).clone() - ) + articulation.set_joint_position_target_index(target=articulation.data.default_joint_pos.torch.clone()) articulation.write_data_to_sim() # perform step sim.step() @@ -1328,7 +1312,7 @@ def test_external_force_on_multiple_bodies_at_position(sim, num_articulations, d # check condition for i in range(num_articulations): # since there is a moment applied on the articulation, the articulation should rotate - assert torch.abs(wp.to_torch(articulation.data.root_ang_vel_w)[i, 2]).item() > 0.1 + assert torch.abs(articulation.data.root_ang_vel_w.torch[i, 2]).item() > 0.1 @pytest.mark.isaacsim_ci @@ -1511,7 +1495,7 @@ def test_setting_velocity_limit_implicit( articulation.root_view.get_attribute("joint_velocity_limit", SimulationManager.get_model()) ).to(device)[:, 0, :] # check data buffer - torch.testing.assert_close(wp.to_torch(articulation.data.joint_velocity_limits), newton_vel_limit) + torch.testing.assert_close(articulation.data.joint_velocity_limits.torch, newton_vel_limit) # check actuator has simulation velocity limit torch.testing.assert_close(articulation.actuators["joint"].velocity_limit_sim, newton_vel_limit) # check that both values match for velocity limit @@ -1566,7 +1550,7 @@ def test_setting_velocity_limit_explicit(sim, num_articulations, device, vel_lim actuator_vel_limit_sim = articulation.actuators["joint"].velocity_limit_sim # check data buffer for joint_velocity_limits_sim - torch.testing.assert_close(wp.to_torch(articulation.data.joint_velocity_limits), newton_vel_limit) + torch.testing.assert_close(articulation.data.joint_velocity_limits.torch, newton_vel_limit) # check actuator velocity_limit_sim is set to physx torch.testing.assert_close(actuator_vel_limit_sim, newton_vel_limit) @@ -1740,10 +1724,10 @@ def test_reset(sim, num_articulations, device, articulation_type): # Reset should zero external forces and torques assert not articulation._instantaneous_wrench_composer.active assert not articulation._permanent_wrench_composer.active - assert torch.count_nonzero(wp.to_torch(articulation._instantaneous_wrench_composer.composed_force)) == 0 - assert torch.count_nonzero(wp.to_torch(articulation._instantaneous_wrench_composer.composed_torque)) == 0 - assert torch.count_nonzero(wp.to_torch(articulation._permanent_wrench_composer.composed_force)) == 0 - assert torch.count_nonzero(wp.to_torch(articulation._permanent_wrench_composer.composed_torque)) == 0 + assert torch.count_nonzero(articulation._instantaneous_wrench_composer.composed_force.torch) == 0 + assert torch.count_nonzero(articulation._instantaneous_wrench_composer.composed_torque.torch) == 0 + assert torch.count_nonzero(articulation._permanent_wrench_composer.composed_force.torch) == 0 + assert torch.count_nonzero(articulation._permanent_wrench_composer.composed_torque.torch) == 0 if num_articulations > 1: num_bodies = articulation.num_bodies @@ -1758,20 +1742,10 @@ def test_reset(sim, num_articulations, device, articulation_type): articulation.reset(env_ids=torch.tensor([0], device=device)) assert articulation._instantaneous_wrench_composer.active assert articulation._permanent_wrench_composer.active - assert ( - torch.count_nonzero(wp.to_torch(articulation._instantaneous_wrench_composer.composed_force)) - == num_bodies * 3 - ) - assert ( - torch.count_nonzero(wp.to_torch(articulation._instantaneous_wrench_composer.composed_torque)) - == num_bodies * 3 - ) - assert ( - torch.count_nonzero(wp.to_torch(articulation._permanent_wrench_composer.composed_force)) == num_bodies * 3 - ) - assert ( - torch.count_nonzero(wp.to_torch(articulation._permanent_wrench_composer.composed_torque)) == num_bodies * 3 - ) + assert torch.count_nonzero(articulation._instantaneous_wrench_composer.composed_force.torch) == num_bodies * 3 + assert torch.count_nonzero(articulation._instantaneous_wrench_composer.composed_torque.torch) == num_bodies * 3 + assert torch.count_nonzero(articulation._permanent_wrench_composer.composed_force.torch) == num_bodies * 3 + assert torch.count_nonzero(articulation._permanent_wrench_composer.composed_torque.torch) == num_bodies * 3 @pytest.mark.isaacsim_ci @@ -1796,7 +1770,7 @@ def test_apply_joint_command(sim, num_articulations, device, add_ground_plane, a articulation.update(sim.cfg.dt) # reset dof state - joint_pos = wp.to_torch(articulation.data.default_joint_pos).clone() + joint_pos = articulation.data.default_joint_pos.torch.clone() joint_pos[:, 3] = 0.0 # apply action to the articulation @@ -1812,7 +1786,7 @@ def test_apply_joint_command(sim, num_articulations, device, add_ground_plane, a # Check that current joint position is not the same as default joint position, meaning # the articulation moved. We can't check that it reached its desired joint position as the gains # are not properly tuned - assert not torch.allclose(wp.to_torch(articulation.data.joint_pos), joint_pos) + assert not torch.allclose(articulation.data.joint_pos.torch, joint_pos) @pytest.mark.isaacsim_ci @@ -1878,19 +1852,19 @@ def test_body_root_state(sim, num_articulations, device, with_offset, articulati articulation.update(sim.cfg.dt) # get state properties - root_link_pose_w = wp.to_torch(articulation.data.root_link_pose_w) - root_link_vel_w = wp.to_torch(articulation.data.root_link_vel_w) - root_com_pose_w = wp.to_torch(articulation.data.root_com_pose_w) - root_com_vel_w = wp.to_torch(articulation.data.root_com_vel_w) - body_link_pose_w = wp.to_torch(articulation.data.body_link_pose_w) - body_link_vel_w = wp.to_torch(articulation.data.body_link_vel_w) - body_com_pose_w = wp.to_torch(articulation.data.body_com_pose_w) - body_com_vel_w = wp.to_torch(articulation.data.body_com_vel_w) + root_link_pose_w = articulation.data.root_link_pose_w.torch + root_link_vel_w = articulation.data.root_link_vel_w.torch + root_com_pose_w = articulation.data.root_com_pose_w.torch + root_com_vel_w = articulation.data.root_com_vel_w.torch + body_link_pose_w = articulation.data.body_link_pose_w.torch + body_link_vel_w = articulation.data.body_link_vel_w.torch + body_com_pose_w = articulation.data.body_com_pose_w.torch + body_com_vel_w = articulation.data.body_com_vel_w.torch if with_offset: # get joint state - joint_pos = wp.to_torch(articulation.data.joint_pos).unsqueeze(-1) - joint_vel = wp.to_torch(articulation.data.joint_vel).unsqueeze(-1) + joint_pos = articulation.data.joint_pos.torch.unsqueeze(-1) + joint_vel = articulation.data.joint_vel.torch.unsqueeze(-1) # LINK state # angular velocity should be the same for both COM and link frames @@ -1925,7 +1899,7 @@ def test_body_root_state(sim, num_articulations, device, with_offset, articulati torch.testing.assert_close(pos_gt, body_com_pose_w[..., :3], atol=1e-3, rtol=1e-1) # orientation - com_quat_b = wp.to_torch(articulation.data.body_com_quat_b) + com_quat_b = articulation.data.body_com_quat_b.torch com_quat_w = math_utils.quat_mul(body_link_pose_w[..., 3:], com_quat_b) torch.testing.assert_close(com_quat_w, body_com_pose_w[..., 3:]) torch.testing.assert_close(com_quat_w[:, root_idx, :], root_com_pose_w[..., 3:]) @@ -1997,7 +1971,7 @@ def test_write_root_state( ) rand_state = torch.zeros(num_articulations, 13, device=device) - rand_state[..., :7] = wp.to_torch(articulation.data.default_root_pose) + rand_state[..., :7] = articulation.data.default_root_pose.torch rand_state[..., :3] += env_pos # make quaternion a unit vector rand_state[..., 3:7] = torch.nn.functional.normalize(rand_state[..., 3:7], dim=-1) @@ -2025,11 +1999,11 @@ def test_write_root_state( articulation.write_root_link_velocity_to_sim_index(root_velocity=rand_state[..., 7:], env_ids=env_idx) if state_location == "com": - torch.testing.assert_close(rand_state[..., :7], wp.to_torch(articulation.data.root_com_pose_w)) - torch.testing.assert_close(rand_state[..., 7:], wp.to_torch(articulation.data.root_com_vel_w)) + torch.testing.assert_close(rand_state[..., :7], articulation.data.root_com_pose_w.torch) + torch.testing.assert_close(rand_state[..., 7:], articulation.data.root_com_vel_w.torch) elif state_location == "link": - torch.testing.assert_close(rand_state[..., :7], wp.to_torch(articulation.data.root_link_pose_w)) - torch.testing.assert_close(rand_state[..., 7:], wp.to_torch(articulation.data.root_link_vel_w)) + torch.testing.assert_close(rand_state[..., :7], articulation.data.root_link_pose_w.torch) + torch.testing.assert_close(rand_state[..., 7:], articulation.data.root_link_vel_w.torch) @pytest.mark.isaacsim_ci @@ -2068,12 +2042,12 @@ def test_body_incoming_joint_wrench_b_single_joint(sim, num_articulations, devic external_torque_vector_b[:, arm_idx, 2] = 10.0 # 10 Nm in z direction # apply action to the articulation - joint_pos = torch.ones_like(wp.to_torch(articulation.data.joint_pos)) * 1.5708 / 2.0 + joint_pos = torch.ones_like(articulation.data.joint_pos.torch) * 1.5708 / 2.0 articulation.write_joint_position_to_sim_index( - position=torch.ones_like(wp.to_torch(articulation.data.joint_pos)), + position=torch.ones_like(articulation.data.joint_pos.torch), ) articulation.write_joint_velocity_to_sim_index( - velocity=torch.zeros_like(wp.to_torch(articulation.data.joint_vel)), + velocity=torch.zeros_like(articulation.data.joint_vel.torch), ) articulation.set_joint_position_target_index(target=joint_pos) articulation.write_data_to_sim() @@ -2088,16 +2062,16 @@ def test_body_incoming_joint_wrench_b_single_joint(sim, num_articulations, devic articulation.update(sim.cfg.dt) # check shape - assert wp.to_torch(articulation.data.body_incoming_joint_wrench_b).shape == ( + assert articulation.data.body_incoming_joint_wrench_b.torch.shape == ( num_articulations, articulation.num_bodies, 6, ) # calculate expected static - mass = wp.to_torch(articulation.data.body_mass).to("cpu") - pos_w = wp.to_torch(articulation.data.body_pos_w) - quat_w = wp.to_torch(articulation.data.body_quat_w) + mass = articulation.data.body_mass.torch.to("cpu") + pos_w = articulation.data.body_pos_w.torch + quat_w = articulation.data.body_quat_w.torch mass_link2 = mass[:, arm_idx].view(num_articulations, -1) gravity = torch.tensor(sim.cfg.gravity, device="cpu").repeat(num_articulations, 1).view((num_articulations, 3)) @@ -2128,7 +2102,7 @@ def test_body_incoming_joint_wrench_b_single_joint(sim, num_articulations, devic # check value of last joint wrench torch.testing.assert_close( expected_wrench, - wp.to_torch(articulation.data.body_incoming_joint_wrench_b)[:, arm_idx, :].squeeze(1), + articulation.data.body_incoming_joint_wrench_b.torch[:, arm_idx, :].squeeze(1), atol=1e-2, rtol=1e-3, ) @@ -2204,13 +2178,13 @@ def test_write_joint_state_data_consistency(sim, num_articulations, device, grav from torch.distributions import Uniform - joint_pos_limits = wp.to_torch(articulation.data.joint_pos_limits) - joint_vel_limits = wp.to_torch(articulation.data.joint_vel_limits) + joint_pos_limits = articulation.data.joint_pos_limits.torch + joint_vel_limits = articulation.data.joint_vel_limits.torch pos_dist = Uniform(joint_pos_limits[..., 0], joint_pos_limits[..., 1]) vel_dist = Uniform(-joint_vel_limits, joint_vel_limits) - original_body_link_pose_w = wp.to_torch(articulation.data.body_link_pose_w).clone() - original_body_com_vel_w = wp.to_torch(articulation.data.body_com_vel_w).clone() + original_body_link_pose_w = articulation.data.body_link_pose_w.torch.clone() + original_body_com_vel_w = articulation.data.body_com_vel_w.torch.clone() rand_joint_pos = pos_dist.sample() rand_joint_vel = vel_dist.sample() @@ -2218,65 +2192,65 @@ def test_write_joint_state_data_consistency(sim, num_articulations, device, grav articulation.write_joint_position_to_sim_index(position=rand_joint_pos) articulation.write_joint_velocity_to_sim_index(velocity=rand_joint_vel) # make sure valued updated - body_link_pose_w = wp.to_torch(articulation.data.body_link_pose_w) - body_com_vel_w = wp.to_torch(articulation.data.body_com_vel_w) + body_link_pose_w = articulation.data.body_link_pose_w.torch + body_com_vel_w = articulation.data.body_com_vel_w.torch original_body_states = torch.cat([original_body_link_pose_w, original_body_com_vel_w], dim=-1) body_state_w = torch.cat([body_link_pose_w, body_com_vel_w], dim=-1) assert torch.count_nonzero(original_body_states[:, 1:] != body_state_w[:, 1:]) > ( len(original_body_states[:, 1:]) / 2 ) # validate body - link consistency - body_link_vel_w = wp.to_torch(articulation.data.body_link_vel_w) - torch.testing.assert_close(body_link_pose_w, wp.to_torch(articulation.data.body_link_pose_w)) + body_link_vel_w = articulation.data.body_link_vel_w.torch + torch.testing.assert_close(body_link_pose_w, articulation.data.body_link_pose_w.torch) # skip lin_vel because it differs from link frame, this should be fine because we are only checking # if velocity update is triggered, which can be determined by comparing angular velocity torch.testing.assert_close(body_com_vel_w[..., 3:], body_link_vel_w[..., 3:]) # validate link - com conistency - body_com_pos_b = wp.to_torch(articulation.data.body_com_pos_b) - body_com_quat_b = wp.to_torch(articulation.data.body_com_quat_b) + body_com_pos_b = articulation.data.body_com_pos_b.torch + body_com_quat_b = articulation.data.body_com_quat_b.torch expected_com_pos, expected_com_quat = math_utils.combine_frame_transforms( body_link_pose_w[..., :3].view(-1, 3), body_link_pose_w[..., 3:].view(-1, 4), body_com_pos_b.view(-1, 3), body_com_quat_b.view(-1, 4), ) - body_com_pos_w = wp.to_torch(articulation.data.body_com_pos_w) - body_com_quat_w = wp.to_torch(articulation.data.body_com_quat_w) + body_com_pos_w = articulation.data.body_com_pos_w.torch + body_com_quat_w = articulation.data.body_com_quat_w.torch torch.testing.assert_close(expected_com_pos.view(len(env_idx), -1, 3), body_com_pos_w) torch.testing.assert_close(expected_com_quat.view(len(env_idx), -1, 4), body_com_quat_w) # validate body - com consistency - body_com_lin_vel_w = wp.to_torch(articulation.data.body_com_lin_vel_w) - body_com_ang_vel_w = wp.to_torch(articulation.data.body_com_ang_vel_w) + body_com_lin_vel_w = articulation.data.body_com_lin_vel_w.torch + body_com_ang_vel_w = articulation.data.body_com_ang_vel_w.torch torch.testing.assert_close(body_com_vel_w[..., :3], body_com_lin_vel_w) torch.testing.assert_close(body_com_vel_w[..., 3:], body_com_ang_vel_w) # validate pos_w, quat_w, pos_b, quat_b is consistent with pose_w and pose_b expected_com_pose_w = torch.cat((body_com_pos_w, body_com_quat_w), dim=2) expected_com_pose_b = torch.cat((body_com_pos_b, body_com_quat_b), dim=2) - body_pos_w = wp.to_torch(articulation.data.body_pos_w) - body_quat_w = wp.to_torch(articulation.data.body_quat_w) + body_pos_w = articulation.data.body_pos_w.torch + body_quat_w = articulation.data.body_quat_w.torch expected_body_pose_w = torch.cat((body_pos_w, body_quat_w), dim=2) - body_link_pos_w = wp.to_torch(articulation.data.body_link_pos_w) - body_link_quat_w = wp.to_torch(articulation.data.body_link_quat_w) + body_link_pos_w = articulation.data.body_link_pos_w.torch + body_link_quat_w = articulation.data.body_link_quat_w.torch expected_body_link_pose_w = torch.cat((body_link_pos_w, body_link_quat_w), dim=2) - body_com_pose_w = wp.to_torch(articulation.data.body_com_pose_w) - body_com_pose_b = wp.to_torch(articulation.data.body_com_pose_b) - body_pose_w = wp.to_torch(articulation.data.body_pose_w) - body_link_pose_w_fresh = wp.to_torch(articulation.data.body_link_pose_w) + body_com_pose_w = articulation.data.body_com_pose_w.torch + body_com_pose_b = articulation.data.body_com_pose_b.torch + body_pose_w = articulation.data.body_pose_w.torch + body_link_pose_w_fresh = articulation.data.body_link_pose_w.torch torch.testing.assert_close(body_com_pose_w, expected_com_pose_w) torch.testing.assert_close(body_com_pose_b, expected_com_pose_b) torch.testing.assert_close(body_pose_w, expected_body_pose_w) torch.testing.assert_close(body_link_pose_w_fresh, expected_body_link_pose_w) # validate pose_w is consistent with individual properties - body_vel_w = wp.to_torch(articulation.data.body_vel_w) - body_com_vel_w_fresh = wp.to_torch(articulation.data.body_com_vel_w) + body_vel_w = articulation.data.body_vel_w.torch + body_com_vel_w_fresh = articulation.data.body_com_vel_w.torch torch.testing.assert_close(body_pose_w, body_link_pose_w) torch.testing.assert_close(body_vel_w, body_com_vel_w) torch.testing.assert_close(body_link_pose_w_fresh, body_link_pose_w) - torch.testing.assert_close(body_com_pose_w, wp.to_torch(articulation.data.body_com_pose_w)) + torch.testing.assert_close(body_com_pose_w, articulation.data.body_com_pose_w.torch) torch.testing.assert_close(body_vel_w, body_com_vel_w_fresh) @@ -2313,11 +2287,11 @@ def test_spatial_tendons(sim, num_articulations, device, articulation_type): # Check that fixed base assert articulation.is_fixed_base # Check buffers that exists and have correct shapes - assert wp.to_torch(articulation.data.root_pos_w).shape == (num_articulations, 3) - assert wp.to_torch(articulation.data.root_quat_w).shape == (num_articulations, 4) - assert wp.to_torch(articulation.data.joint_pos).shape == (num_articulations, 3) - assert wp.to_torch(articulation.data.body_mass).shape == (num_articulations, articulation.num_bodies) - assert wp.to_torch(articulation.data.body_inertia).shape == (num_articulations, articulation.num_bodies, 9) + assert articulation.data.root_pos_w.torch.shape == (num_articulations, 3) + assert articulation.data.root_quat_w.torch.shape == (num_articulations, 4) + assert articulation.data.joint_pos.torch.shape == (num_articulations, 3) + assert articulation.data.body_mass.torch.shape == (num_articulations, articulation.num_bodies) + assert articulation.data.body_inertia.torch.shape == (num_articulations, articulation.num_bodies, 9) assert articulation.num_spatial_tendons == 1 articulation.set_spatial_tendon_stiffness_index(stiffness=10.0) @@ -2448,7 +2422,7 @@ def test_body_q_consistent_after_root_write(num_articulations, device, articulat articulation.update(sim.cfg.dt) # Teleport env 0 by 10m (simulating a reset) - new_pose = wp.to_torch(articulation.data.default_root_pose).clone() + new_pose = articulation.data.default_root_pose.torch.clone() new_pose[0, 0] += 10.0 new_pose[0, 1] += 5.0 articulation.write_root_pose_to_sim_index( @@ -2555,7 +2529,7 @@ def test_randomize_rigid_body_com(sim, num_articulations, device, add_ground_pla sim.reset() assert articulation.is_initialized - original_com = wp.to_torch(articulation.data.body_com_pos_b).clone() + original_com = articulation.data.body_com_pos_b.torch.clone() com_offset = torch.zeros(num_articulations, articulation.num_bodies, 3, device=device) com_offset[..., 0] = 0.5 @@ -2563,7 +2537,7 @@ def test_randomize_rigid_body_com(sim, num_articulations, device, add_ground_pla env_ids = torch.arange(num_articulations, device=device, dtype=torch.int32) articulation.set_coms_index(coms=new_com, env_ids=env_ids) - updated_com = wp.to_torch(articulation.data.body_com_pos_b) + updated_com = articulation.data.body_com_pos_b.torch torch.testing.assert_close(updated_com, new_com, atol=1e-5, rtol=1e-5) diff --git a/source/isaaclab_newton/test/assets/test_rigid_object.py b/source/isaaclab_newton/test/assets/test_rigid_object.py index f08eb4d76d80..3b8e3aa9bf85 100644 --- a/source/isaaclab_newton/test/assets/test_rigid_object.py +++ b/source/isaaclab_newton/test/assets/test_rigid_object.py @@ -140,10 +140,10 @@ def test_initialization(num_cubes, device): assert len(cube_object.body_names) == 1 # Check buffers that exists and have correct shapes - assert wp.to_torch(cube_object.data.root_pos_w).shape == (num_cubes, 3) - assert wp.to_torch(cube_object.data.root_quat_w).shape == (num_cubes, 4) - assert wp.to_torch(cube_object.data.body_mass).shape == (num_cubes, 1) - assert wp.to_torch(cube_object.data.body_inertia).shape == (num_cubes, 1, 9) + assert cube_object.data.root_pos_w.torch.shape == (num_cubes, 3) + assert cube_object.data.root_quat_w.torch.shape == (num_cubes, 4) + assert cube_object.data.body_mass.torch.shape == (num_cubes, 1) + assert cube_object.data.body_inertia.torch.shape == (num_cubes, 1, 9) # Simulate physics for _ in range(2): @@ -175,8 +175,8 @@ def test_initialization_with_kinematic_enabled(num_cubes, device): assert len(cube_object.body_names) == 1 # Check buffers that exists and have correct shapes - assert wp.to_torch(cube_object.data.root_pos_w).shape == (num_cubes, 3) - assert wp.to_torch(cube_object.data.root_quat_w).shape == (num_cubes, 4) + assert cube_object.data.root_pos_w.torch.shape == (num_cubes, 3) + assert cube_object.data.root_quat_w.torch.shape == (num_cubes, 4) # Simulate physics for _ in range(2): @@ -185,11 +185,11 @@ def test_initialization_with_kinematic_enabled(num_cubes, device): # update object cube_object.update(sim.cfg.dt) # check that the object is kinematic - default_root_pose = wp.to_torch(cube_object.data.default_root_pose).clone() - default_root_vel = wp.to_torch(cube_object.data.default_root_vel).clone() + default_root_pose = cube_object.data.default_root_pose.torch.clone() + default_root_vel = cube_object.data.default_root_vel.torch.clone() default_root_pose[:, :3] += origins - torch.testing.assert_close(wp.to_torch(cube_object.data.root_link_pose_w), default_root_pose) - torch.testing.assert_close(wp.to_torch(cube_object.data.root_com_vel_w), default_root_vel) + torch.testing.assert_close(cube_object.data.root_link_pose_w.torch, default_root_pose) + torch.testing.assert_close(cube_object.data.root_com_vel_w.torch, default_root_vel) @pytest.mark.isaacsim_ci @@ -276,8 +276,8 @@ def test_external_force_buffer(device): # check if the cube's force and torque buffers are correctly updated for i in range(cube_object.num_instances): - assert wp.to_torch(cube_object._permanent_wrench_composer.composed_force)[i, 0, 0].item() == force - assert wp.to_torch(cube_object._permanent_wrench_composer.composed_torque)[i, 0, 0].item() == force + assert cube_object._permanent_wrench_composer.composed_force.torch[i, 0, 0].item() == force + assert cube_object._permanent_wrench_composer.composed_torque.torch[i, 0, 0].item() == force # Check if the instantaneous wrench is correctly added to the permanent wrench cube_object.permanent_wrench_composer.add_forces_and_torques_index( @@ -322,13 +322,13 @@ def test_external_force_on_single_body(num_cubes, device): # Sample a force equal to the weight of the object external_wrench_b = torch.zeros(cube_object.num_instances, len(body_ids), 6, device=sim.device) # Every 2nd cube should have a force applied to it - external_wrench_b[0::2, :, 2] = 9.81 * wp.to_torch(cube_object.data.body_mass)[0] + external_wrench_b[0::2, :, 2] = 9.81 * cube_object.data.body_mass.torch[0] # Now we are ready! for i in range(5): # reset root state - root_pose = wp.to_torch(cube_object.data.default_root_pose).clone() - root_vel = wp.to_torch(cube_object.data.default_root_vel).clone() + root_pose = cube_object.data.default_root_pose.torch.clone() + root_vel = cube_object.data.default_root_vel.torch.clone() # need to shift the position of the cubes otherwise they will be on top of each other root_pose[:, :3] = origins @@ -341,7 +341,7 @@ def test_external_force_on_single_body(num_cubes, device): is_global = False if i % 2 == 0: is_global = True - positions = wp.to_torch(cube_object.data.body_com_pos_w)[:, body_ids, :3] + positions = cube_object.data.body_com_pos_w.torch[:, body_ids, :3] else: positions = None @@ -366,10 +366,10 @@ def test_external_force_on_single_body(num_cubes, device): # First object should still be at the same Z position (1.0) torch.testing.assert_close( - wp.to_torch(cube_object.data.root_pos_w)[0::2, 2], torch.ones(num_cubes // 2, device=sim.device) + cube_object.data.root_pos_w.torch[0::2, 2], torch.ones(num_cubes // 2, device=sim.device) ) # Second object should have fallen, so it's Z height should be less than initial height of 1.0 - assert torch.all(wp.to_torch(cube_object.data.root_pos_w)[1::2, 2] < 1.0) + assert torch.all(cube_object.data.root_pos_w.torch[1::2, 2] < 1.0) @pytest.mark.parametrize("num_cubes", [2, 4]) @@ -404,8 +404,8 @@ def test_external_force_on_single_body_at_position(num_cubes, device): # Now we are ready! for i in range(5): # reset root state - root_pose = wp.to_torch(cube_object.data.default_root_pose).clone() - root_vel = wp.to_torch(cube_object.data.default_root_vel).clone() + root_pose = cube_object.data.default_root_pose.torch.clone() + root_vel = cube_object.data.default_root_vel.torch.clone() # need to shift the position of the cubes otherwise they will be on top of each other root_pose[:, :3] = origins @@ -418,7 +418,7 @@ def test_external_force_on_single_body_at_position(num_cubes, device): is_global = False if i % 2 == 0: is_global = True - body_com_pos_w = wp.to_torch(cube_object.data.body_com_pos_w)[:, body_ids, :3] + body_com_pos_w = cube_object.data.body_com_pos_w.torch[:, body_ids, :3] external_wrench_positions_b[..., 0] = 0.0 external_wrench_positions_b[..., 1] = 1.0 external_wrench_positions_b[..., 2] = 0.0 @@ -455,9 +455,9 @@ def test_external_force_on_single_body_at_position(num_cubes, device): cube_object.update(sim.cfg.dt) # The first object should be rotating around it's X axis - assert torch.all(torch.abs(wp.to_torch(cube_object.data.root_ang_vel_b)[0::2, 0]) > 0.1) + assert torch.all(torch.abs(cube_object.data.root_ang_vel_b.torch[0::2, 0]) > 0.1) # Second object should have fallen, so it's Z height should be less than initial height of 1.0 - assert torch.all(wp.to_torch(cube_object.data.root_pos_w)[1::2, 2] < 1.0) + assert torch.all(cube_object.data.root_pos_w.torch[1::2, 2] < 1.0) @pytest.mark.isaacsim_ci @@ -485,10 +485,10 @@ def test_set_rigid_object_state(num_cubes, device): # Set each state type individually as they are dependent on each other for state_type_to_randomize in state_types: state_dict = { - "root_pos_w": torch.zeros_like(wp.to_torch(cube_object.data.root_pos_w), device=sim.device), + "root_pos_w": torch.zeros_like(cube_object.data.root_pos_w.torch, device=sim.device), "root_quat_w": default_orientation(num=num_cubes, device=sim.device), - "root_lin_vel_w": torch.zeros_like(wp.to_torch(cube_object.data.root_lin_vel_w), device=sim.device), - "root_ang_vel_w": torch.zeros_like(wp.to_torch(cube_object.data.root_ang_vel_w), device=sim.device), + "root_lin_vel_w": torch.zeros_like(cube_object.data.root_lin_vel_w.torch, device=sim.device), + "root_ang_vel_w": torch.zeros_like(cube_object.data.root_ang_vel_w.torch, device=sim.device), } # Now we are ready! @@ -520,7 +520,7 @@ def test_set_rigid_object_state(num_cubes, device): # assert that set root quantities are equal to the ones set in the state_dict for key, expected_value in state_dict.items(): - value = wp.to_torch(getattr(cube_object.data, key)) + value = getattr(cube_object.data, key).torch # Newton reads state directly from sim (not cached), so post-step drift # from velocity integration causes larger differences than PhysX torch.testing.assert_close(value, expected_value, rtol=1e-1, atol=1e-1) @@ -549,13 +549,13 @@ def test_reset_rigid_object(num_cubes, device): cube_object.update(sim.cfg.dt) # Move the object to a random position - root_pose = wp.to_torch(cube_object.data.default_root_pose).clone() + root_pose = cube_object.data.default_root_pose.torch.clone() root_pose[:, :3] = torch.randn(num_cubes, 3, device=sim.device) # Random orientation root_pose[:, 3:7] = random_orientation(num=num_cubes, device=sim.device) cube_object.write_root_pose_to_sim_index(root_pose=root_pose) - root_vel = wp.to_torch(cube_object.data.default_root_vel).clone() + root_vel = cube_object.data.default_root_vel.torch.clone() cube_object.write_root_velocity_to_sim_index(root_velocity=root_vel) if i % 2 == 0: @@ -565,10 +565,10 @@ def test_reset_rigid_object(num_cubes, device): # Reset should zero external forces and torques assert not cube_object._instantaneous_wrench_composer.active assert not cube_object._permanent_wrench_composer.active - assert torch.count_nonzero(wp.to_torch(cube_object._instantaneous_wrench_composer.composed_force)) == 0 - assert torch.count_nonzero(wp.to_torch(cube_object._instantaneous_wrench_composer.composed_torque)) == 0 - assert torch.count_nonzero(wp.to_torch(cube_object._permanent_wrench_composer.composed_force)) == 0 - assert torch.count_nonzero(wp.to_torch(cube_object._permanent_wrench_composer.composed_torque)) == 0 + assert torch.count_nonzero(cube_object._instantaneous_wrench_composer.composed_force.torch) == 0 + assert torch.count_nonzero(cube_object._instantaneous_wrench_composer.composed_torque.torch) == 0 + assert torch.count_nonzero(cube_object._permanent_wrench_composer.composed_force.torch) == 0 + assert torch.count_nonzero(cube_object._permanent_wrench_composer.composed_torque.torch) == 0 @pytest.mark.isaacsim_ci @@ -674,7 +674,7 @@ def test_rigid_body_no_friction(num_cubes, device): tolerance = 1e-5 torch.testing.assert_close( - wp.to_torch(cube_object.data.root_lin_vel_w), initial_velocity[:, :3], rtol=1e-5, atol=tolerance + cube_object.data.root_lin_vel_w.torch, initial_velocity[:, :3], rtol=1e-5, atol=tolerance ) @@ -720,7 +720,7 @@ def test_rigid_body_with_static_friction(num_cubes, device): sim.step() cube_object.update(sim.cfg.dt) cube_object.write_root_velocity_to_sim_index(root_velocity=torch.zeros((num_cubes, 6), device=sim.device)) - cube_mass = wp.to_torch(cube_object.data.body_mass) + cube_mass = cube_object.data.body_mass.torch gravity_magnitude = abs(sim.cfg.gravity[2]) # 2 cases: force applied is below and above mu # below mu: block should not move as the force applied is <= mu @@ -741,7 +741,7 @@ def test_rigid_body_with_static_friction(num_cubes, device): ) # Get root state - initial_root_pos = wp.to_torch(cube_object.data.root_pos_w).clone() + initial_root_pos = cube_object.data.root_pos_w.torch.clone() # Simulate physics for _ in range(200): # apply the wrench @@ -752,10 +752,10 @@ def test_rigid_body_with_static_friction(num_cubes, device): if force == "below_mu": # Assert that the block has not moved torch.testing.assert_close( - wp.to_torch(cube_object.data.root_pos_w), initial_root_pos, rtol=2e-3, atol=2e-3 + cube_object.data.root_pos_w.torch, initial_root_pos, rtol=2e-3, atol=2e-3 ) if force == "above_mu": - assert (wp.to_torch(cube_object.data.root_pos_w)[..., 0] - initial_root_pos[..., 0] > 0.02).all() + assert (cube_object.data.root_pos_w.torch[..., 0] - initial_root_pos[..., 0] > 0.02).all() @pytest.mark.isaacsim_ci @@ -811,14 +811,14 @@ def test_rigid_body_with_restitution(num_cubes, device): device=device, ) - curr_z_velocity = wp.to_torch(cube_object.data.root_lin_vel_w)[:, 2].clone() + curr_z_velocity = cube_object.data.root_lin_vel_w.torch[:, 2].clone() for _ in range(100): sim.step() # update object cube_object.update(sim.cfg.dt) - curr_z_velocity = wp.to_torch(cube_object.data.root_lin_vel_w)[:, 2].clone() + curr_z_velocity = cube_object.data.root_lin_vel_w.torch[:, 2].clone() if expected_collision_type == "inelastic": # assert that the block has not bounced by checking that the z velocity is less than or equal to 0 @@ -851,7 +851,7 @@ def test_rigid_body_set_mass(num_cubes, device): sim.reset() # Get masses before increasing - original_masses = wp.to_torch(cube_object.data.body_mass) + original_masses = cube_object.data.body_mass.torch assert original_masses.shape == (num_cubes, 1) @@ -861,7 +861,7 @@ def test_rigid_body_set_mass(num_cubes, device): # Set masses using Newton API cube_object.set_masses_index(masses=wp.from_torch(masses, dtype=wp.float32)) - torch.testing.assert_close(wp.to_torch(cube_object.data.body_mass), masses) + torch.testing.assert_close(cube_object.data.body_mass.torch, masses) # Simulate physics # perform rendering @@ -869,7 +869,7 @@ def test_rigid_body_set_mass(num_cubes, device): # update object cube_object.update(sim.cfg.dt) - masses_to_check = wp.to_torch(cube_object.data.body_mass) + masses_to_check = cube_object.data.body_mass.torch # Check if mass is set correctly torch.testing.assert_close(masses, masses_to_check) @@ -896,9 +896,9 @@ def test_gravity_vec_w(num_cubes, device, gravity_enabled): sim.reset() # Check that gravity is set correctly - assert wp.to_torch(cube_object.data.GRAVITY_VEC_W)[0, 0] == gravity_dir[0] - assert wp.to_torch(cube_object.data.GRAVITY_VEC_W)[0, 1] == gravity_dir[1] - assert wp.to_torch(cube_object.data.GRAVITY_VEC_W)[0, 2] == gravity_dir[2] + assert cube_object.data.GRAVITY_VEC_W.torch[0, 0] == gravity_dir[0] + assert cube_object.data.GRAVITY_VEC_W.torch[0, 1] == gravity_dir[1] + assert cube_object.data.GRAVITY_VEC_W.torch[0, 2] == gravity_dir[2] # Simulate physics for _ in range(2): @@ -912,7 +912,7 @@ def test_gravity_vec_w(num_cubes, device, gravity_enabled): if gravity_enabled: gravity[:, :, 2] = -9.81 # Check the body accelerations are correct - torch.testing.assert_close(wp.to_torch(cube_object.data.body_acc_w), gravity) + torch.testing.assert_close(cube_object.data.body_acc_w.torch, gravity) @pytest.mark.isaacsim_ci @@ -946,7 +946,7 @@ def test_body_root_state_properties(num_cubes, device, with_offset): SimulationManager._solver.notify_model_changed(SolverNotifyFlags.BODY_INERTIAL_PROPERTIES) # check center of mass has been set - torch.testing.assert_close(wp.to_torch(cube_object.data.body_com_pos_b).squeeze(1), offset) + torch.testing.assert_close(cube_object.data.body_com_pos_b.torch.squeeze(1), offset) # random z spin velocity (bounded to keep numerical drift within the position tolerance below) spin_twist = torch.zeros(6, device=device) @@ -962,14 +962,14 @@ def test_body_root_state_properties(num_cubes, device, with_offset): cube_object.update(sim.cfg.dt) # get state properties - root_link_pose_w = wp.to_torch(cube_object.data.root_link_pose_w) - root_link_vel_w = wp.to_torch(cube_object.data.root_link_vel_w) - root_com_pose_w = wp.to_torch(cube_object.data.root_com_pose_w) - root_com_vel_w = wp.to_torch(cube_object.data.root_com_vel_w) - body_link_pose_w = wp.to_torch(cube_object.data.body_link_pose_w) - body_link_vel_w = wp.to_torch(cube_object.data.body_link_vel_w) - body_com_pose_w = wp.to_torch(cube_object.data.body_com_pose_w) - body_com_vel_w = wp.to_torch(cube_object.data.body_com_vel_w) + root_link_pose_w = cube_object.data.root_link_pose_w.torch + root_link_vel_w = cube_object.data.root_link_vel_w.torch + root_com_pose_w = cube_object.data.root_com_pose_w.torch + root_com_vel_w = cube_object.data.root_com_vel_w.torch + body_link_pose_w = cube_object.data.body_link_pose_w.torch + body_link_vel_w = cube_object.data.body_link_vel_w.torch + body_com_pose_w = cube_object.data.body_com_pose_w.torch + body_com_vel_w = cube_object.data.body_com_vel_w.torch # if offset is [0,0,0] all root_state_%_w will match and all body_%_w will match if not with_offset: @@ -1001,7 +1001,7 @@ def test_body_root_state_properties(num_cubes, device, with_offset): torch.testing.assert_close(-offset, body_link_state_pos_rel_com.squeeze(-2), **_tol) # orientation of com will be a constant rotation from link orientation - com_quat_b = wp.to_torch(cube_object.data.body_com_quat_b) + com_quat_b = cube_object.data.body_com_quat_b.torch com_quat_w = quat_mul(body_link_pose_w[..., 3:], com_quat_b) torch.testing.assert_close(com_quat_w, body_com_pose_w[..., 3:], **_tol) torch.testing.assert_close(com_quat_w.squeeze(-2), root_com_pose_w[..., 3:], **_tol) @@ -1060,10 +1060,10 @@ def test_write_root_state(num_cubes, device, with_offset, state_location): SimulationManager._solver.notify_model_changed(SolverNotifyFlags.BODY_INERTIAL_PROPERTIES) # check center of mass has been set - torch.testing.assert_close(wp.to_torch(cube_object.data.body_com_pos_b).squeeze(1), offset) + torch.testing.assert_close(cube_object.data.body_com_pos_b.torch.squeeze(1), offset) rand_state = torch.zeros(num_cubes, 13, device=device) - rand_state[..., :7] = wp.to_torch(cube_object.data.default_root_pose) + rand_state[..., :7] = cube_object.data.default_root_pose.torch rand_state[..., :3] += env_pos # make quaternion a unit vector rand_state[..., 3:7] = torch.nn.functional.normalize(rand_state[..., 3:7], dim=-1) @@ -1092,11 +1092,11 @@ def test_write_root_state(num_cubes, device, with_offset, state_location): ) if state_location == "com": - torch.testing.assert_close(rand_state[..., :7], wp.to_torch(cube_object.data.root_com_pose_w)) - torch.testing.assert_close(rand_state[..., 7:], wp.to_torch(cube_object.data.root_com_vel_w)) + torch.testing.assert_close(rand_state[..., :7], cube_object.data.root_com_pose_w.torch) + torch.testing.assert_close(rand_state[..., 7:], cube_object.data.root_com_vel_w.torch) elif state_location == "link": - torch.testing.assert_close(rand_state[..., :7], wp.to_torch(cube_object.data.root_link_pose_w)) - torch.testing.assert_close(rand_state[..., 7:], wp.to_torch(cube_object.data.root_link_vel_w)) + torch.testing.assert_close(rand_state[..., :7], cube_object.data.root_link_pose_w.torch) + torch.testing.assert_close(rand_state[..., 7:], cube_object.data.root_link_vel_w.torch) @pytest.mark.isaacsim_ci @@ -1130,7 +1130,7 @@ def test_write_state_functions_data_consistency(num_cubes, device, with_offset, SimulationManager._solver.notify_model_changed(SolverNotifyFlags.BODY_INERTIAL_PROPERTIES) # check center of mass has been set - torch.testing.assert_close(wp.to_torch(cube_object.data.body_com_pos_b).squeeze(1), offset) + torch.testing.assert_close(cube_object.data.body_com_pos_b.torch.squeeze(1), offset) rand_state = torch.rand(num_cubes, 13, device=device) rand_state[..., :3] += env_pos @@ -1153,9 +1153,9 @@ def test_write_state_functions_data_consistency(num_cubes, device, with_offset, cube_object.write_root_velocity_to_sim_index(root_velocity=rand_state[..., 7:]) if state_location == "com": - root_com_pose_w = wp.to_torch(cube_object.data.root_com_pose_w) - root_com_vel_w = wp.to_torch(cube_object.data.root_com_vel_w) - body_com_pose_b = wp.to_torch(cube_object.data.body_com_pose_b) + root_com_pose_w = cube_object.data.root_com_pose_w.torch + root_com_vel_w = cube_object.data.root_com_vel_w.torch + body_com_pose_b = cube_object.data.body_com_pose_b.torch expected_root_link_pos, expected_root_link_quat = combine_frame_transforms( root_com_pose_w[:, :3], root_com_pose_w[:, 3:], @@ -1163,19 +1163,19 @@ def test_write_state_functions_data_consistency(num_cubes, device, with_offset, quat_inv(body_com_pose_b[:, 0, 3:7]), ) expected_root_link_pose = torch.cat((expected_root_link_pos, expected_root_link_quat), dim=1) - root_link_pose_w = wp.to_torch(cube_object.data.root_link_pose_w) - root_link_vel_w = wp.to_torch(cube_object.data.root_link_vel_w) + root_link_pose_w = cube_object.data.root_link_pose_w.torch + root_link_vel_w = cube_object.data.root_link_vel_w.torch # test both root_pose and root_link successfully updated when root_com updates torch.testing.assert_close(expected_root_link_pose, root_link_pose_w) # skip lin_vel because it differs from link frame, this should be fine because we are only checking # if velocity update is triggered, which can be determined by comparing angular velocity torch.testing.assert_close(root_com_vel_w[:, 3:], root_link_vel_w[:, 3:]) torch.testing.assert_close(expected_root_link_pose, root_link_pose_w) - torch.testing.assert_close(root_com_vel_w[:, 3:], wp.to_torch(cube_object.data.root_com_vel_w)[:, 3:]) + torch.testing.assert_close(root_com_vel_w[:, 3:], cube_object.data.root_com_vel_w.torch[:, 3:]) elif state_location == "link": - root_link_pose_w = wp.to_torch(cube_object.data.root_link_pose_w) - root_link_vel_w = wp.to_torch(cube_object.data.root_link_vel_w) - body_com_pose_b = wp.to_torch(cube_object.data.body_com_pose_b) + root_link_pose_w = cube_object.data.root_link_pose_w.torch + root_link_vel_w = cube_object.data.root_link_vel_w.torch + body_com_pose_b = cube_object.data.body_com_pose_b.torch expected_com_pos, expected_com_quat = combine_frame_transforms( root_link_pose_w[:, :3], root_link_pose_w[:, 3:], @@ -1183,19 +1183,19 @@ def test_write_state_functions_data_consistency(num_cubes, device, with_offset, body_com_pose_b[:, 0, 3:7], ) expected_com_pose = torch.cat((expected_com_pos, expected_com_quat), dim=1) - root_com_pose_w = wp.to_torch(cube_object.data.root_com_pose_w) - root_com_vel_w = wp.to_torch(cube_object.data.root_com_vel_w) + root_com_pose_w = cube_object.data.root_com_pose_w.torch + root_com_vel_w = cube_object.data.root_com_vel_w.torch # test both root_pose and root_com successfully updated when root_link updates torch.testing.assert_close(expected_com_pose, root_com_pose_w) # skip lin_vel because it differs from link frame, this should be fine because we are only checking # if velocity update is triggered, which can be determined by comparing angular velocity torch.testing.assert_close(root_link_vel_w[:, 3:], root_com_vel_w[:, 3:]) - torch.testing.assert_close(root_link_pose_w, wp.to_torch(cube_object.data.root_link_pose_w)) - torch.testing.assert_close(root_link_vel_w[:, 3:], wp.to_torch(cube_object.data.root_com_vel_w)[:, 3:]) + torch.testing.assert_close(root_link_pose_w, cube_object.data.root_link_pose_w.torch) + torch.testing.assert_close(root_link_vel_w[:, 3:], cube_object.data.root_com_vel_w.torch[:, 3:]) elif state_location == "root": - root_link_pose_w = wp.to_torch(cube_object.data.root_link_pose_w) - root_com_vel_w = wp.to_torch(cube_object.data.root_com_vel_w) - body_com_pose_b = wp.to_torch(cube_object.data.body_com_pose_b) + root_link_pose_w = cube_object.data.root_link_pose_w.torch + root_com_vel_w = cube_object.data.root_com_vel_w.torch + body_com_pose_b = cube_object.data.body_com_pose_b.torch expected_com_pos, expected_com_quat = combine_frame_transforms( root_link_pose_w[:, :3], root_link_pose_w[:, 3:], @@ -1203,12 +1203,12 @@ def test_write_state_functions_data_consistency(num_cubes, device, with_offset, body_com_pose_b[:, 0, 3:7], ) expected_com_pose = torch.cat((expected_com_pos, expected_com_quat), dim=1) - root_com_pose_w = wp.to_torch(cube_object.data.root_com_pose_w) - root_link_vel_w = wp.to_torch(cube_object.data.root_link_vel_w) + root_com_pose_w = cube_object.data.root_com_pose_w.torch + root_link_vel_w = cube_object.data.root_link_vel_w.torch # test both root_com and root_link successfully updated when root_pose updates torch.testing.assert_close(expected_com_pose, root_com_pose_w) - torch.testing.assert_close(root_com_vel_w, wp.to_torch(cube_object.data.root_com_vel_w)) - torch.testing.assert_close(root_link_pose_w, wp.to_torch(cube_object.data.root_link_pose_w)) + torch.testing.assert_close(root_com_vel_w, cube_object.data.root_com_vel_w.torch) + torch.testing.assert_close(root_link_pose_w, cube_object.data.root_link_pose_w.torch) torch.testing.assert_close(root_com_vel_w[:, 3:], root_link_vel_w[:, 3:]) diff --git a/source/isaaclab_newton/test/assets/test_rigid_object_collection.py b/source/isaaclab_newton/test/assets/test_rigid_object_collection.py index 4c5599e35887..0873da8ecf22 100644 --- a/source/isaaclab_newton/test/assets/test_rigid_object_collection.py +++ b/source/isaaclab_newton/test/assets/test_rigid_object_collection.py @@ -137,10 +137,10 @@ def test_initialization(num_envs, num_cubes, device): assert len(object_collection.body_names) == num_cubes # Check buffers that exist and have correct shapes - assert wp.to_torch(object_collection.data.body_link_pos_w).shape == (num_envs, num_cubes, 3) - assert wp.to_torch(object_collection.data.body_link_quat_w).shape == (num_envs, num_cubes, 4) - assert wp.to_torch(object_collection.data.body_mass).shape == (num_envs, num_cubes) - assert wp.to_torch(object_collection.data.body_inertia).shape == (num_envs, num_cubes, 9) + assert object_collection.data.body_link_pos_w.torch.shape == (num_envs, num_cubes, 3) + assert object_collection.data.body_link_quat_w.torch.shape == (num_envs, num_cubes, 4) + assert object_collection.data.body_mass.torch.shape == (num_envs, num_cubes) + assert object_collection.data.body_inertia.torch.shape == (num_envs, num_cubes, 9) # Simulate physics for _ in range(2): @@ -171,19 +171,19 @@ def test_initialization_with_kinematic_enabled(num_envs, num_cubes, device): assert len(object_collection.body_names) == num_cubes # Check buffers that exist and have correct shapes - assert wp.to_torch(object_collection.data.body_link_pos_w).shape == (num_envs, num_cubes, 3) - assert wp.to_torch(object_collection.data.body_link_quat_w).shape == (num_envs, num_cubes, 4) + assert object_collection.data.body_link_pos_w.torch.shape == (num_envs, num_cubes, 3) + assert object_collection.data.body_link_quat_w.torch.shape == (num_envs, num_cubes, 4) # Simulate physics for _ in range(2): sim.step() object_collection.update(sim.cfg.dt) # check that the object is kinematic - default_body_pose = wp.to_torch(object_collection.data.default_body_pose).clone() - default_body_vel = wp.to_torch(object_collection.data.default_body_vel).clone() + default_body_pose = object_collection.data.default_body_pose.torch.clone() + default_body_vel = object_collection.data.default_body_vel.torch.clone() default_body_pose[..., :3] += origins.unsqueeze(1) - torch.testing.assert_close(wp.to_torch(object_collection.data.body_link_pose_w), default_body_pose) - torch.testing.assert_close(wp.to_torch(object_collection.data.body_link_vel_w), default_body_vel) + torch.testing.assert_close(object_collection.data.body_link_pose_w.torch, default_body_pose) + torch.testing.assert_close(object_collection.data.body_link_vel_w.torch, default_body_vel) @pytest.mark.parametrize("num_cubes", [1, 2]) @@ -241,10 +241,8 @@ def test_external_force_buffer(device): # check if the object collection's force and torque buffers are correctly updated for i in range(num_envs): - assert wp.to_torch(object_collection._permanent_wrench_composer.composed_force)[i, 0, 0].item() == force - assert ( - wp.to_torch(object_collection._permanent_wrench_composer.composed_torque)[i, 0, 0].item() == force - ) + assert object_collection._permanent_wrench_composer.composed_force.torch[i, 0, 0].item() == force + assert object_collection._permanent_wrench_composer.composed_torque.torch[i, 0, 0].item() == force object_collection.instantaneous_wrench_composer.add_forces_and_torques_index( body_ids=object_ids, @@ -274,12 +272,12 @@ def test_external_force_on_single_body(num_envs, num_cubes, device): # Sample a force equal to the weight of the object external_wrench_b = torch.zeros(object_collection.num_instances, len(object_ids), 6, device=sim.device) # Every 2nd cube should have a force applied to it - external_wrench_b[:, 0::2, 2] = 9.81 * wp.to_torch(object_collection.data.body_mass)[:, 0::2] + external_wrench_b[:, 0::2, 2] = 9.81 * object_collection.data.body_mass.torch[:, 0::2] for i in range(5): # reset object state - body_pose = wp.to_torch(object_collection.data.default_body_pose).clone() - body_vel = wp.to_torch(object_collection.data.default_body_vel).clone() + body_pose = object_collection.data.default_body_pose.torch.clone() + body_vel = object_collection.data.default_body_vel.torch.clone() # need to shift the position of the cubes otherwise they will be on top of each other body_pose[..., :2] += origins.unsqueeze(1)[..., :2] object_collection.write_body_link_pose_to_sim_index(body_poses=body_pose) @@ -289,7 +287,7 @@ def test_external_force_on_single_body(num_envs, num_cubes, device): is_global = False if i % 2 == 0: - positions = wp.to_torch(object_collection.data.body_link_pos_w)[:, object_ids, :3] + positions = object_collection.data.body_link_pos_w.torch[:, object_ids, :3] is_global = True else: positions = None @@ -313,11 +311,11 @@ def test_external_force_on_single_body(num_envs, num_cubes, device): # First object should still be at the same Z position (1.0) torch.testing.assert_close( - wp.to_torch(object_collection.data.body_link_pos_w)[:, 0::2, 2], - torch.ones_like(wp.to_torch(object_collection.data.body_link_pos_w)[:, 0::2, 2]), + object_collection.data.body_link_pos_w.torch[:, 0::2, 2], + torch.ones_like(object_collection.data.body_link_pos_w.torch[:, 0::2, 2]), ) # Second object should have fallen, so it's Z height should be less than initial height of 1.0 - assert torch.all(wp.to_torch(object_collection.data.body_link_pos_w)[:, 1::2, 2] < 1.0) + assert torch.all(object_collection.data.body_link_pos_w.torch[:, 1::2, 2] < 1.0) @pytest.mark.parametrize("num_envs", [1, 2]) @@ -350,8 +348,8 @@ def test_external_force_on_single_body_at_position(num_envs, num_cubes, device): # Desired force and torque for i in range(5): # reset object state - body_pose = wp.to_torch(object_collection.data.default_body_pose).clone() - body_vel = wp.to_torch(object_collection.data.default_body_vel).clone() + body_pose = object_collection.data.default_body_pose.torch.clone() + body_vel = object_collection.data.default_body_vel.torch.clone() # need to shift the position of the cubes otherwise they will be on top of each other body_pose[..., :2] += origins.unsqueeze(1)[..., :2] object_collection.write_body_link_pose_to_sim_index(body_poses=body_pose) @@ -361,7 +359,7 @@ def test_external_force_on_single_body_at_position(num_envs, num_cubes, device): is_global = False if i % 2 == 0: - body_com_pos_w = wp.to_torch(object_collection.data.body_link_pos_w)[:, object_ids, :3] + body_com_pos_w = object_collection.data.body_link_pos_w.torch[:, object_ids, :3] external_wrench_positions_b[..., 0] = 0.0 external_wrench_positions_b[..., 1] = 1.0 external_wrench_positions_b[..., 2] = 0.0 @@ -398,9 +396,9 @@ def test_external_force_on_single_body_at_position(num_envs, num_cubes, device): object_collection.update(sim.cfg.dt) # First object should be rotating around it's X axis - assert torch.all(wp.to_torch(object_collection.data.body_com_ang_vel_b)[:, 0::2, 0] > 0.1) + assert torch.all(object_collection.data.body_com_ang_vel_b.torch[:, 0::2, 0] > 0.1) # Second object should have fallen, so it's Z height should be less than initial height of 1.0 - assert torch.all(wp.to_torch(object_collection.data.body_link_pos_w)[:, 1::2, 2] < 1.0) + assert torch.all(object_collection.data.body_link_pos_w.torch[:, 1::2, 2] < 1.0) @pytest.mark.parametrize("num_envs", [1, 3]) @@ -423,17 +421,15 @@ def test_set_object_state(num_envs, num_cubes, device): # Set each state type individually as they are dependent on each other for state_type_to_randomize in state_types: state_dict = { - "body_link_pos_w": torch.zeros_like( - wp.to_torch(object_collection.data.body_link_pos_w), device=sim.device - ), + "body_link_pos_w": torch.zeros_like(object_collection.data.body_link_pos_w.torch, device=sim.device), "body_link_quat_w": default_orientation(num=num_cubes * num_envs, device=sim.device).view( num_envs, num_cubes, 4 ), "body_com_lin_vel_w": torch.zeros_like( - wp.to_torch(object_collection.data.body_com_lin_vel_w), device=sim.device + object_collection.data.body_com_lin_vel_w.torch, device=sim.device ), "body_com_ang_vel_w": torch.zeros_like( - wp.to_torch(object_collection.data.body_com_ang_vel_w), device=sim.device + object_collection.data.body_com_ang_vel_w.torch, device=sim.device ), } @@ -469,7 +465,7 @@ def test_set_object_state(num_envs, num_cubes, device): # assert that set object quantities are equal to the ones set in the state_dict for key, expected_value in state_dict.items(): - value = wp.to_torch(getattr(object_collection.data, key)) + value = getattr(object_collection.data, key).torch # Newton reads state directly from sim (not cached), so post-step drift # from velocity integration causes larger differences than PhysX torch.testing.assert_close(value, expected_value, rtol=1e-1, atol=1e-1) @@ -492,12 +488,12 @@ def test_reset_object_collection(num_envs, num_cubes, device): object_collection.update(sim.cfg.dt) # Move the object to a random position - body_pose = wp.to_torch(object_collection.data.default_body_pose).clone() + body_pose = object_collection.data.default_body_pose.torch.clone() body_pose[..., :3] = torch.randn(num_envs, num_cubes, 3, device=sim.device) # Random orientation body_pose[..., 3:7] = random_orientation(num=num_cubes, device=sim.device) object_collection.write_body_link_pose_to_sim_index(body_poses=body_pose) - body_vel = wp.to_torch(object_collection.data.default_body_vel).clone() + body_vel = object_collection.data.default_body_vel.torch.clone() object_collection.write_body_com_velocity_to_sim_index(body_velocities=body_vel) if i % 2 == 0: @@ -506,20 +502,10 @@ def test_reset_object_collection(num_envs, num_cubes, device): # Reset should zero external forces and torques assert not object_collection._instantaneous_wrench_composer.active assert not object_collection._permanent_wrench_composer.active - assert ( - torch.count_nonzero(wp.to_torch(object_collection._instantaneous_wrench_composer.composed_force)) - == 0 - ) - assert ( - torch.count_nonzero(wp.to_torch(object_collection._instantaneous_wrench_composer.composed_torque)) - == 0 - ) - assert ( - torch.count_nonzero(wp.to_torch(object_collection._permanent_wrench_composer.composed_force)) == 0 - ) - assert ( - torch.count_nonzero(wp.to_torch(object_collection._permanent_wrench_composer.composed_torque)) == 0 - ) + assert torch.count_nonzero(object_collection._instantaneous_wrench_composer.composed_force.torch) == 0 + assert torch.count_nonzero(object_collection._instantaneous_wrench_composer.composed_torque.torch) == 0 + assert torch.count_nonzero(object_collection._permanent_wrench_composer.composed_force.torch) == 0 + assert torch.count_nonzero(object_collection._permanent_wrench_composer.composed_torque.torch) == 0 @pytest.mark.parametrize("num_envs", [1, 3]) @@ -577,7 +563,7 @@ def test_gravity_vec_w(num_envs, num_cubes, device, gravity_enabled): sim.reset() # Check if gravity vector is set correctly - gravity_vec = wp.to_torch(object_collection.data.GRAVITY_VEC_W) + gravity_vec = object_collection.data.GRAVITY_VEC_W.torch assert gravity_vec[0, 0, 0] == gravity_dir[0] assert gravity_vec[0, 0, 1] == gravity_dir[1] assert gravity_vec[0, 0, 2] == gravity_dir[2] @@ -593,7 +579,7 @@ def test_gravity_vec_w(num_envs, num_cubes, device, gravity_enabled): gravity[..., 2] = -9.81 # Check the body accelerations are correct - torch.testing.assert_close(wp.to_torch(object_collection.data.body_com_acc_w), gravity) + torch.testing.assert_close(object_collection.data.body_com_acc_w.torch, gravity) @pytest.mark.parametrize("num_envs", [1, 4]) @@ -625,14 +611,14 @@ def test_object_state_properties(num_envs, num_cubes, device, with_offset): SimulationManager._solver.notify_model_changed(SolverNotifyFlags.BODY_INERTIAL_PROPERTIES) # check center of mass has been set - torch.testing.assert_close(wp.to_torch(cube_object.data.body_com_pos_b), offset) + torch.testing.assert_close(cube_object.data.body_com_pos_b.torch, offset) # random z spin velocity spin_twist = torch.zeros(6, device=device) spin_twist[5] = torch.randn(1, device=device) # initial spawn point - init_com = wp.to_torch(cube_object.data.body_com_pose_w)[..., :3] + init_com = cube_object.data.body_com_pose_w.torch[..., :3] for i in range(10): # spin the object around Z axis (com) @@ -641,10 +627,10 @@ def test_object_state_properties(num_envs, num_cubes, device, with_offset): cube_object.update(sim.cfg.dt) # get state properties - object_link_pose_w = wp.to_torch(cube_object.data.body_link_pose_w) - object_link_vel_w = wp.to_torch(cube_object.data.body_link_vel_w) - object_com_pose_w = wp.to_torch(cube_object.data.body_com_pose_w) - object_com_vel_w = wp.to_torch(cube_object.data.body_com_vel_w) + object_link_pose_w = cube_object.data.body_link_pose_w.torch + object_link_vel_w = cube_object.data.body_link_vel_w.torch + object_com_pose_w = cube_object.data.body_com_pose_w.torch + object_com_vel_w = cube_object.data.body_com_vel_w.torch # if offset is [0,0,0] all object_state_%_w will match and all body_%_w will match if not with_offset: @@ -666,7 +652,7 @@ def test_object_state_properties(num_envs, num_cubes, device, with_offset): torch.testing.assert_close(-offset, object_link_state_pos_rel_com, **_tol) # orientation of com will be a constant rotation from link orientation - com_quat_b = wp.to_torch(cube_object.data.body_com_quat_b) + com_quat_b = cube_object.data.body_com_quat_b.torch com_quat_w = quat_mul(object_link_pose_w[..., 3:], com_quat_b) torch.testing.assert_close(com_quat_w, object_com_pose_w[..., 3:], **_tol) @@ -724,11 +710,11 @@ def test_write_object_state(num_envs, num_cubes, device, with_offset, state_loca SimulationManager._solver.notify_model_changed(SolverNotifyFlags.BODY_INERTIAL_PROPERTIES) # check center of mass has been set - torch.testing.assert_close(wp.to_torch(cube_object.data.body_com_pos_b), offset) + torch.testing.assert_close(cube_object.data.body_com_pos_b.torch, offset) rand_state = torch.zeros(num_envs, num_cubes, 13, device=device) - rand_state[..., :7] = wp.to_torch(cube_object.data.default_body_pose) - rand_state[..., :3] += wp.to_torch(cube_object.data.body_link_pos_w) + rand_state[..., :7] = cube_object.data.default_body_pose.torch + rand_state[..., :3] += cube_object.data.body_link_pos_w.torch # make quaternion a unit vector rand_state[..., 3:7] = torch.nn.functional.normalize(rand_state[..., 3:7], dim=-1) @@ -762,11 +748,11 @@ def test_write_object_state(num_envs, num_cubes, device, with_offset, state_loca ) if state_location == "com": - torch.testing.assert_close(rand_state[..., :7], wp.to_torch(cube_object.data.body_com_pose_w)) - torch.testing.assert_close(rand_state[..., 7:], wp.to_torch(cube_object.data.body_com_vel_w)) + torch.testing.assert_close(rand_state[..., :7], cube_object.data.body_com_pose_w.torch) + torch.testing.assert_close(rand_state[..., 7:], cube_object.data.body_com_vel_w.torch) elif state_location == "link": - torch.testing.assert_close(rand_state[..., :7], wp.to_torch(cube_object.data.body_link_pose_w)) - torch.testing.assert_close(rand_state[..., 7:], wp.to_torch(cube_object.data.body_link_vel_w)) + torch.testing.assert_close(rand_state[..., :7], cube_object.data.body_link_pose_w.torch) + torch.testing.assert_close(rand_state[..., 7:], cube_object.data.body_link_vel_w.torch) @pytest.mark.parametrize("num_envs", [1, 3]) @@ -802,10 +788,10 @@ def test_write_object_state_functions_data_consistency(num_envs, num_cubes, devi SimulationManager._solver.notify_model_changed(SolverNotifyFlags.BODY_INERTIAL_PROPERTIES) # check center of mass has been set - torch.testing.assert_close(wp.to_torch(cube_object.data.body_com_pos_b), offset) + torch.testing.assert_close(cube_object.data.body_com_pos_b.torch, offset) rand_state = torch.rand(num_envs, num_cubes, 13, device=device) - rand_state[..., :3] += wp.to_torch(cube_object.data.body_link_pos_w) + rand_state[..., :3] += cube_object.data.body_link_pos_w.torch # make quaternion a unit vector rand_state[..., 3:7] = torch.nn.functional.normalize(rand_state[..., 3:7], dim=-1) @@ -814,8 +800,8 @@ def test_write_object_state_functions_data_consistency(num_envs, num_cubes, devi sim.step() cube_object.update(sim.cfg.dt) - body_link_pose_w = wp.to_torch(cube_object.data.body_link_pose_w) - body_com_pose_w = wp.to_torch(cube_object.data.body_com_pose_w) + body_link_pose_w = cube_object.data.body_link_pose_w.torch + body_com_pose_w = cube_object.data.body_com_pose_w.torch object_link_to_com_pos, object_link_to_com_quat = subtract_frame_transforms( body_link_pose_w[..., :3].view(-1, 3), body_link_pose_w[..., 3:7].view(-1, 4), @@ -846,8 +832,8 @@ def test_write_object_state_functions_data_consistency(num_envs, num_cubes, devi ) if state_location == "com": - com_pose_w = wp.to_torch(cube_object.data.body_com_pose_w) - com_vel_w = wp.to_torch(cube_object.data.body_com_vel_w) + com_pose_w = cube_object.data.body_com_pose_w.torch + com_vel_w = cube_object.data.body_com_vel_w.torch expected_root_link_pos, expected_root_link_quat = combine_frame_transforms( com_pose_w[..., :3].view(-1, 3), com_pose_w[..., 3:].view(-1, 4), @@ -857,18 +843,18 @@ def test_write_object_state_functions_data_consistency(num_envs, num_cubes, devi expected_object_link_pose = torch.cat((expected_root_link_pos, expected_root_link_quat), dim=1).view( num_envs, -1, 7 ) - link_pose_w = wp.to_torch(cube_object.data.body_link_pose_w) - link_vel_w = wp.to_torch(cube_object.data.body_link_vel_w) + link_pose_w = cube_object.data.body_link_pose_w.torch + link_vel_w = cube_object.data.body_link_vel_w.torch # test both root_pose and root_link successfully updated when root_com updates torch.testing.assert_close(expected_object_link_pose, link_pose_w) # skip lin_vel because it differs from link frame, this should be fine because we are only checking # if velocity update is triggered, which can be determined by comparing angular velocity torch.testing.assert_close(com_vel_w[..., 3:], link_vel_w[..., 3:]) torch.testing.assert_close(expected_object_link_pose, link_pose_w) - torch.testing.assert_close(com_vel_w[..., 3:], wp.to_torch(cube_object.data.body_com_vel_w)[..., 3:]) + torch.testing.assert_close(com_vel_w[..., 3:], cube_object.data.body_com_vel_w.torch[..., 3:]) elif state_location == "link": - link_pose_w = wp.to_torch(cube_object.data.body_link_pose_w) - link_vel_w = wp.to_torch(cube_object.data.body_link_vel_w) + link_pose_w = cube_object.data.body_link_pose_w.torch + link_vel_w = cube_object.data.body_link_vel_w.torch expected_com_pos, expected_com_quat = combine_frame_transforms( link_pose_w[..., :3].view(-1, 3), link_pose_w[..., 3:].view(-1, 4), @@ -876,18 +862,18 @@ def test_write_object_state_functions_data_consistency(num_envs, num_cubes, devi object_link_to_com_quat, ) expected_object_com_pose = torch.cat((expected_com_pos, expected_com_quat), dim=1).view(num_envs, -1, 7) - com_pose_w = wp.to_torch(cube_object.data.body_com_pose_w) - com_vel_w = wp.to_torch(cube_object.data.body_com_vel_w) + com_pose_w = cube_object.data.body_com_pose_w.torch + com_vel_w = cube_object.data.body_com_vel_w.torch # test both root_pose and root_com successfully updated when root_link updates torch.testing.assert_close(expected_object_com_pose, com_pose_w) # skip lin_vel because it differs from link frame, this should be fine because we are only checking # if velocity update is triggered, which can be determined by comparing angular velocity torch.testing.assert_close(link_vel_w[..., 3:], com_vel_w[..., 3:]) - torch.testing.assert_close(link_pose_w, wp.to_torch(cube_object.data.body_link_pose_w)) - torch.testing.assert_close(link_vel_w[..., 3:], wp.to_torch(cube_object.data.body_com_vel_w)[..., 3:]) + torch.testing.assert_close(link_pose_w, cube_object.data.body_link_pose_w.torch) + torch.testing.assert_close(link_vel_w[..., 3:], cube_object.data.body_com_vel_w.torch[..., 3:]) elif state_location == "root": - body_link_pose_w = wp.to_torch(cube_object.data.body_link_pose_w) - body_com_vel_w = wp.to_torch(cube_object.data.body_com_vel_w) + body_link_pose_w = cube_object.data.body_link_pose_w.torch + body_com_vel_w = cube_object.data.body_com_vel_w.torch expected_object_com_pos, expected_object_com_quat = combine_frame_transforms( body_link_pose_w[..., :3].view(-1, 3), body_link_pose_w[..., 3:].view(-1, 4), @@ -897,10 +883,10 @@ def test_write_object_state_functions_data_consistency(num_envs, num_cubes, devi expected_object_com_pose = torch.cat((expected_object_com_pos, expected_object_com_quat), dim=1).view( num_envs, -1, 7 ) - com_pose_w = wp.to_torch(cube_object.data.body_com_pose_w) - com_vel_w = wp.to_torch(cube_object.data.body_com_vel_w) - link_pose_w = wp.to_torch(cube_object.data.body_link_pose_w) - link_vel_w = wp.to_torch(cube_object.data.body_link_vel_w) + com_pose_w = cube_object.data.body_com_pose_w.torch + com_vel_w = cube_object.data.body_com_vel_w.torch + link_pose_w = cube_object.data.body_link_pose_w.torch + link_vel_w = cube_object.data.body_link_vel_w.torch # test both root_com and root_link successfully updated when root_pose updates torch.testing.assert_close(expected_object_com_pose, com_pose_w) torch.testing.assert_close(body_com_vel_w, com_vel_w) diff --git a/source/isaaclab_newton/test/sensors/test_contact_sensor.py b/source/isaaclab_newton/test/sensors/test_contact_sensor.py index f5da322df4b7..066803ee884b 100644 --- a/source/isaaclab_newton/test/sensors/test_contact_sensor.py +++ b/source/isaaclab_newton/test/sensors/test_contact_sensor.py @@ -26,7 +26,6 @@ import pytest import torch -import warp as wp from physics.physics_test_utils import ( COLLISION_PIPELINES, STABLE_SHAPES, @@ -124,7 +123,7 @@ def test_contact_lifecycle(device: str, use_mujoco_contacts: bool, shape_type: S contact_sensor: ContactSensor = scene["contact_sensor_a"] obj: RigidObject = scene["object_a"] - root_pose = wp.to_torch(obj.data.root_link_pose_w).clone() + root_pose = obj.data.root_link_pose_w.torch.clone() for group_idx, base_height in enumerate(base_heights): for i in range(envs_per_group): env_idx = group_idx * envs_per_group + i @@ -139,13 +138,13 @@ def test_contact_lifecycle(device: str, use_mujoco_contacts: bool, shape_type: S for _ in range(5): perform_sim_step(sim, scene, SIM_DT) - forces = torch.norm(wp.to_torch(contact_sensor.data.net_forces_w), dim=-1) + forces = torch.norm(contact_sensor.data.net_forces_w.torch, dim=-1) for env_idx in range(num_envs): assert forces[env_idx].max().item() < 0.01, f"Env {env_idx}: No contact should be detected while in air." for tick in range(5, total_fall_steps): perform_sim_step(sim, scene, SIM_DT) - forces = torch.norm(wp.to_torch(contact_sensor.data.net_forces_w), dim=-1) + forces = torch.norm(contact_sensor.data.net_forces_w.torch, dim=-1) for env_idx in range(num_envs): if forces[env_idx].max().item() > 0.1 and not contact_detected[env_idx]: contact_detected[env_idx] = True @@ -184,7 +183,7 @@ def test_contact_lifecycle(device: str, use_mujoco_contacts: bool, shape_type: S for step in range(lift_steps): perform_sim_step(sim, scene, SIM_DT) if step > 10: - forces = torch.norm(wp.to_torch(contact_sensor.data.net_forces_w), dim=-1) + forces = torch.norm(contact_sensor.data.net_forces_w.torch, dim=-1) for env_idx in range(num_envs): if forces[env_idx].max().item() < 0.01: no_contact_detected[env_idx] = True @@ -257,8 +256,8 @@ def test_horizontal_collision_detects_contact(device: str, use_mujoco_contacts: sensor_a: ContactSensor = scene["contact_sensor_a"] sensor_b: ContactSensor = scene["contact_sensor_b"] - pose_a = wp.to_torch(object_a.data.root_link_pose_w).clone() - pose_b = wp.to_torch(object_b.data.root_link_pose_w).clone() + pose_a = object_a.data.root_link_pose_w.torch.clone() + pose_b = object_b.data.root_link_pose_w.torch.clone() for group_idx, (_, separation) in enumerate(group_configs): for i in range(envs_per_group): env_idx = group_idx * envs_per_group + i @@ -278,8 +277,8 @@ def test_horizontal_collision_detects_contact(device: str, use_mujoco_contacts: for tick in range(collision_steps): perform_sim_step(sim, scene, SIM_DT) - forces_a = torch.norm(wp.to_torch(sensor_a.data.net_forces_w), dim=-1) - forces_b = torch.norm(wp.to_torch(sensor_b.data.net_forces_w), dim=-1) + forces_a = torch.norm(sensor_a.data.net_forces_w.torch, dim=-1) + forces_b = torch.norm(sensor_b.data.net_forces_w.torch, dim=-1) for env_idx in range(num_envs): if forces_a[env_idx].max().item() > 0.1: contact_detected_a[env_idx] = True @@ -371,8 +370,8 @@ def test_resting_object_contact_force(device: str, use_mujoco_contacts: bool): for step in range(settle_steps): perform_sim_step(sim, scene, SIM_DT) if step >= settle_steps - avg_window: - force_a_samples.append(wp.to_torch(sensor_a.data.net_forces_w).clone()) - force_b_samples.append(wp.to_torch(sensor_b.data.net_forces_w).clone()) + force_a_samples.append(sensor_a.data.net_forces_w.torch.clone()) + force_b_samples.append(sensor_b.data.net_forces_w.torch.clone()) forces_a = torch.stack(force_a_samples).mean(dim=0) forces_b = torch.stack(force_b_samples).mean(dim=0) @@ -444,7 +443,7 @@ def test_higher_drop_produces_larger_impact_force(device: str, use_mujoco_contac obj: RigidObject = scene["object_a"] contact_sensor: ContactSensor = scene["contact_sensor_a"] - root_pose = wp.to_torch(obj.data.root_link_pose_w).clone() + root_pose = obj.data.root_link_pose_w.torch.clone() for env_idx in range(num_envs): root_pose[env_idx, 2] = drop_heights[env_idx] + object_radius obj.write_root_pose_to_sim_index(root_pose=root_pose) @@ -455,7 +454,7 @@ def test_higher_drop_produces_larger_impact_force(device: str, use_mujoco_contac for _ in range(total_steps): perform_sim_step(sim, scene, SIM_DT) - force_magnitudes = torch.norm(wp.to_torch(contact_sensor.data.net_forces_w), dim=-1) + force_magnitudes = torch.norm(contact_sensor.data.net_forces_w.torch, dim=-1) for env_idx in range(num_envs): f = force_magnitudes[env_idx].max().item() if f > 0.1: @@ -572,8 +571,8 @@ def test_filter_enables_force_matrix(device: str, use_mujoco_contacts: bool): net_raw = contact_sensor.data.net_forces_w if not matrix_samples: assert matrix_raw is not None, "force_matrix_w should not be None when filter is set" - matrix_samples.append(wp.to_torch(matrix_raw).clone()) - net_samples.append(wp.to_torch(net_raw).clone()) + matrix_samples.append(matrix_raw.torch.clone()) + net_samples.append(net_raw.torch.clone()) force_matrix = torch.stack(matrix_samples).mean(dim=0) net_forces = torch.stack(net_samples).mean(dim=0) @@ -723,14 +722,14 @@ def test_finger_contact_sensor_isolation(device: str, use_mujoco_contacts: bool, # Newton's articulation.reset() doesn't write default_joint_pos to sim (unlike # ManagerBasedEnv's reset_scene_to_default event). Without this, joints start at 0.0 # which is below thumb_joint_0's lower limit (0.279 rad), causing violent oscillation. - default_jpos = wp.to_torch(hand.data.default_joint_pos).clone() - default_jvel = wp.to_torch(hand.data.default_joint_vel).clone() + default_jpos = hand.data.default_joint_pos.torch.clone() + default_jvel = hand.data.default_joint_vel.torch.clone() hand.write_joint_position_to_sim_index(position=default_jpos) hand.write_joint_velocity_to_sim_index(velocity=default_jvel) hand.set_joint_position_target_index(target=default_jpos) - hand_world_pos = wp.to_torch(hand.data.root_link_pose_w)[:, :3] - drop_pose = wp.to_torch(drop_object.data.root_link_pose_w).clone() + hand_world_pos = hand.data.root_link_pose_w.torch[:, :3] + drop_pose = drop_object.data.root_link_pose_w.torch.clone() for env_idx, finger in enumerate(finger_names): offset = ALLEGRO_FINGERTIP_OFFSETS[finger] drop_pose[env_idx, 0] = hand_world_pos[env_idx, 0] + offset[0] @@ -741,9 +740,9 @@ def test_finger_contact_sensor_isolation(device: str, use_mujoco_contacts: bool, for _ in range(30): perform_sim_step(sim, scene, SIM_DT) - hand_world_pos = wp.to_torch(hand.data.root_link_pose_w)[:, :3] + hand_world_pos = hand.data.root_link_pose_w.torch[:, :3] drop_object.reset() - drop_pose = wp.to_torch(drop_object.data.root_link_pose_w).clone() + drop_pose = drop_object.data.root_link_pose_w.torch.clone() for env_idx, finger in enumerate(finger_names): offset = ALLEGRO_FINGERTIP_OFFSETS[finger] drop_pose[env_idx, 0] = hand_world_pos[env_idx, 0] + offset[0] @@ -761,7 +760,7 @@ def test_finger_contact_sensor_isolation(device: str, use_mujoco_contacts: bool, perform_sim_step(sim, scene, SIM_DT) for finger_name, sensor in finger_sensors.items(): if sensor.data.net_forces_w is not None: - forces = wp.to_torch(sensor.data.net_forces_w) + forces = sensor.data.net_forces_w.torch for env_idx in range(num_envs): f = torch.norm(forces[env_idx]).item() peak_forces[finger_name][env_idx] = max(peak_forces[finger_name][env_idx], f) diff --git a/source/isaaclab_newton/test/sensors/test_frame_transformer.py b/source/isaaclab_newton/test/sensors/test_frame_transformer.py index 453c84f77ba0..513ba9dda918 100644 --- a/source/isaaclab_newton/test/sensors/test_frame_transformer.py +++ b/source/isaaclab_newton/test/sensors/test_frame_transformer.py @@ -15,7 +15,6 @@ import pytest import scipy.spatial.transform as tf import torch -import warp as wp from isaaclab_newton.physics import MJWarpSolverCfg, NewtonCfg import isaaclab.sim as sim_utils @@ -161,7 +160,7 @@ def test_frame_transformer_feet_wrt_base(sim): feet_indices = [feet_indices[i] for i in reordering_indices] # default joint targets - default_actions = wp.to_torch(scene.articulations["robot"].data.default_joint_pos).clone() + default_actions = scene.articulations["robot"].data.default_joint_pos.torch.clone() # Define simulation stepping sim_dt = sim.get_physics_dt() # Simulate physics @@ -169,10 +168,10 @@ def test_frame_transformer_feet_wrt_base(sim): # # reset if count % 25 == 0: # reset root state - root_state = wp.to_torch(scene.articulations["robot"].data.default_root_state).clone() + root_state = scene.articulations["robot"].data.default_root_state.torch.clone() root_state[:, :3] += scene.env_origins - joint_pos = scene.articulations["robot"].data.default_joint_pos - joint_vel = scene.articulations["robot"].data.default_joint_vel + joint_pos = scene.articulations["robot"].data.default_joint_pos.torch + joint_vel = scene.articulations["robot"].data.default_joint_vel.torch # -- set root state # -- robot scene.articulations["robot"].write_root_pose_to_sim(root_state[:, :7]) @@ -193,14 +192,14 @@ def test_frame_transformer_feet_wrt_base(sim): # check absolute frame transforms in world frame # -- ground-truth - root_pose_w = wp.to_torch(scene.articulations["robot"].data.root_pose_w) - feet_pos_w_gt = wp.to_torch(scene.articulations["robot"].data.body_pos_w)[:, feet_indices] - feet_quat_w_gt = wp.to_torch(scene.articulations["robot"].data.body_quat_w)[:, feet_indices] + root_pose_w = scene.articulations["robot"].data.root_pose_w.torch + feet_pos_w_gt = scene.articulations["robot"].data.body_pos_w.torch[:, feet_indices] + feet_quat_w_gt = scene.articulations["robot"].data.body_quat_w.torch[:, feet_indices] # -- frame transformer - source_pos_w_tf = wp.to_torch(scene.sensors["frame_transformer"].data.source_pos_w) - source_quat_w_tf = wp.to_torch(scene.sensors["frame_transformer"].data.source_quat_w) - feet_pos_w_tf = wp.to_torch(scene.sensors["frame_transformer"].data.target_pos_w) - feet_quat_w_tf = wp.to_torch(scene.sensors["frame_transformer"].data.target_quat_w) + source_pos_w_tf = scene.sensors["frame_transformer"].data.source_pos_w.torch + source_quat_w_tf = scene.sensors["frame_transformer"].data.source_quat_w.torch + feet_pos_w_tf = scene.sensors["frame_transformer"].data.target_pos_w.torch + feet_quat_w_tf = scene.sensors["frame_transformer"].data.target_quat_w.torch # check if they are same torch.testing.assert_close(root_pose_w[:, :3], source_pos_w_tf) @@ -209,8 +208,8 @@ def test_frame_transformer_feet_wrt_base(sim): torch.testing.assert_close(feet_quat_w_gt, feet_quat_w_tf) # check if relative transforms are same - feet_pos_source_tf = wp.to_torch(scene.sensors["frame_transformer"].data.target_pos_source) - feet_quat_source_tf = wp.to_torch(scene.sensors["frame_transformer"].data.target_quat_source) + feet_pos_source_tf = scene.sensors["frame_transformer"].data.target_pos_source.torch + feet_quat_source_tf = scene.sensors["frame_transformer"].data.target_quat_source.torch for index in range(len(feet_indices)): # ground-truth foot_pos_b, foot_quat_b = math_utils.subtract_frame_transforms( @@ -259,7 +258,7 @@ def test_frame_transformer_feet_wrt_thigh(sim): assert scene.sensors["frame_transformer"].data.target_frame_names == user_feet_names # default joint targets - default_actions = wp.to_torch(scene.articulations["robot"].data.default_joint_pos).clone() + default_actions = scene.articulations["robot"].data.default_joint_pos.torch.clone() # Define simulation stepping sim_dt = sim.get_physics_dt() # Simulate physics @@ -267,10 +266,10 @@ def test_frame_transformer_feet_wrt_thigh(sim): # # reset if count % 25 == 0: # reset root state - root_state = wp.to_torch(scene.articulations["robot"].data.default_root_state).clone() + root_state = scene.articulations["robot"].data.default_root_state.torch.clone() root_state[:, :3] += scene.env_origins - joint_pos = scene.articulations["robot"].data.default_joint_pos - joint_vel = scene.articulations["robot"].data.default_joint_vel + joint_pos = scene.articulations["robot"].data.default_joint_pos.torch + joint_vel = scene.articulations["robot"].data.default_joint_vel.torch # -- set root state # -- robot scene.articulations["robot"].write_root_pose_to_sim(root_state[:, :7]) @@ -291,14 +290,14 @@ def test_frame_transformer_feet_wrt_thigh(sim): # check absolute frame transforms in world frame # -- ground-truth - source_pose_w_gt = wp.to_torch(scene.articulations["robot"].data.body_state_w)[:, source_frame_index, :7] - feet_pos_w_gt = wp.to_torch(scene.articulations["robot"].data.body_pos_w)[:, feet_indices] - feet_quat_w_gt = wp.to_torch(scene.articulations["robot"].data.body_quat_w)[:, feet_indices] + source_pose_w_gt = scene.articulations["robot"].data.body_state_w.torch[:, source_frame_index, :7] + feet_pos_w_gt = scene.articulations["robot"].data.body_pos_w.torch[:, feet_indices] + feet_quat_w_gt = scene.articulations["robot"].data.body_quat_w.torch[:, feet_indices] # -- frame transformer - source_pos_w_tf = wp.to_torch(scene.sensors["frame_transformer"].data.source_pos_w) - source_quat_w_tf = wp.to_torch(scene.sensors["frame_transformer"].data.source_quat_w) - feet_pos_w_tf = wp.to_torch(scene.sensors["frame_transformer"].data.target_pos_w) - feet_quat_w_tf = wp.to_torch(scene.sensors["frame_transformer"].data.target_quat_w) + source_pos_w_tf = scene.sensors["frame_transformer"].data.source_pos_w.torch + source_quat_w_tf = scene.sensors["frame_transformer"].data.source_quat_w.torch + feet_pos_w_tf = scene.sensors["frame_transformer"].data.target_pos_w.torch + feet_quat_w_tf = scene.sensors["frame_transformer"].data.target_quat_w.torch # check if they are same torch.testing.assert_close(source_pose_w_gt[:, :3], source_pos_w_tf) torch.testing.assert_close(source_pose_w_gt[:, 3:], source_quat_w_tf) @@ -306,8 +305,8 @@ def test_frame_transformer_feet_wrt_thigh(sim): torch.testing.assert_close(feet_quat_w_gt, feet_quat_w_tf) # check if relative transforms are same - feet_pos_source_tf = wp.to_torch(scene.sensors["frame_transformer"].data.target_pos_source) - feet_quat_source_tf = wp.to_torch(scene.sensors["frame_transformer"].data.target_quat_source) + feet_pos_source_tf = scene.sensors["frame_transformer"].data.target_pos_source.torch + feet_quat_source_tf = scene.sensors["frame_transformer"].data.target_quat_source.torch for index in range(len(feet_indices)): # ground-truth foot_pos_b, foot_quat_b = math_utils.subtract_frame_transforms( @@ -337,7 +336,7 @@ def test_frame_transformer_robot_body_to_external_cube(sim): sim.reset() # default joint targets - default_actions = wp.to_torch(scene.articulations["robot"].data.default_joint_pos).clone() + default_actions = scene.articulations["robot"].data.default_joint_pos.torch.clone() # Define simulation stepping sim_dt = sim.get_physics_dt() # Simulate physics @@ -345,10 +344,10 @@ def test_frame_transformer_robot_body_to_external_cube(sim): # # reset if count % 25 == 0: # reset root state - root_state = wp.to_torch(scene.articulations["robot"].data.default_root_state).clone() + root_state = scene.articulations["robot"].data.default_root_state.torch.clone() root_state[:, :3] += scene.env_origins - joint_pos = scene.articulations["robot"].data.default_joint_pos - joint_vel = scene.articulations["robot"].data.default_joint_vel + joint_pos = scene.articulations["robot"].data.default_joint_pos.torch + joint_vel = scene.articulations["robot"].data.default_joint_vel.torch # -- set root state # -- robot scene.articulations["robot"].write_root_pose_to_sim(root_state[:, :7]) @@ -369,14 +368,14 @@ def test_frame_transformer_robot_body_to_external_cube(sim): # check absolute frame transforms in world frame # -- ground-truth - root_pose_w = wp.to_torch(scene.articulations["robot"].data.root_pose_w) - cube_pos_w_gt = wp.to_torch(scene.rigid_objects["cube"].data.root_pos_w) - cube_quat_w_gt = wp.to_torch(scene.rigid_objects["cube"].data.root_quat_w) + root_pose_w = scene.articulations["robot"].data.root_pose_w.torch + cube_pos_w_gt = scene.rigid_objects["cube"].data.root_pos_w.torch + cube_quat_w_gt = scene.rigid_objects["cube"].data.root_quat_w.torch # -- frame transformer - source_pos_w_tf = wp.to_torch(scene.sensors["frame_transformer"].data.source_pos_w) - source_quat_w_tf = wp.to_torch(scene.sensors["frame_transformer"].data.source_quat_w) - cube_pos_w_tf = wp.to_torch(scene.sensors["frame_transformer"].data.target_pos_w).squeeze() - cube_quat_w_tf = wp.to_torch(scene.sensors["frame_transformer"].data.target_quat_w).squeeze() + source_pos_w_tf = scene.sensors["frame_transformer"].data.source_pos_w.torch + source_quat_w_tf = scene.sensors["frame_transformer"].data.source_quat_w.torch + cube_pos_w_tf = scene.sensors["frame_transformer"].data.target_pos_w.torch.squeeze() + cube_quat_w_tf = scene.sensors["frame_transformer"].data.target_quat_w.torch.squeeze() # check if they are same torch.testing.assert_close(root_pose_w[:, :3], source_pos_w_tf) @@ -385,8 +384,8 @@ def test_frame_transformer_robot_body_to_external_cube(sim): torch.testing.assert_close(cube_quat_w_gt, cube_quat_w_tf) # check if relative transforms are same - cube_pos_source_tf = wp.to_torch(scene.sensors["frame_transformer"].data.target_pos_source) - cube_quat_source_tf = wp.to_torch(scene.sensors["frame_transformer"].data.target_quat_source) + cube_pos_source_tf = scene.sensors["frame_transformer"].data.target_pos_source.torch + cube_quat_source_tf = scene.sensors["frame_transformer"].data.target_quat_source.torch # ground-truth cube_pos_b, cube_quat_b = math_utils.subtract_frame_transforms( root_pose_w[:, :3], root_pose_w[:, 3:], cube_pos_w_tf, cube_quat_w_tf @@ -441,7 +440,7 @@ def test_frame_transformer_offset_frames(sim): # # reset if count % 25 == 0: # reset root state - root_state = wp.to_torch(scene["cube"].data.default_root_state).clone() + root_state = scene["cube"].data.default_root_state.torch.clone() root_state[:, :3] += scene.env_origins # -- set root state # -- cube @@ -459,13 +458,13 @@ def test_frame_transformer_offset_frames(sim): # check absolute frame transforms in world frame # -- ground-truth - cube_pos_w_gt = wp.to_torch(scene["cube"].data.root_pos_w) - cube_quat_w_gt = wp.to_torch(scene["cube"].data.root_quat_w) + cube_pos_w_gt = scene["cube"].data.root_pos_w.torch + cube_quat_w_gt = scene["cube"].data.root_quat_w.torch # -- frame transformer - source_pos_w_tf = wp.to_torch(scene.sensors["frame_transformer"].data.source_pos_w) - source_quat_w_tf = wp.to_torch(scene.sensors["frame_transformer"].data.source_quat_w) - target_pos_w_tf = wp.to_torch(scene.sensors["frame_transformer"].data.target_pos_w).squeeze() - target_quat_w_tf = wp.to_torch(scene.sensors["frame_transformer"].data.target_quat_w).squeeze() + source_pos_w_tf = scene.sensors["frame_transformer"].data.source_pos_w.torch + source_quat_w_tf = scene.sensors["frame_transformer"].data.source_quat_w.torch + target_pos_w_tf = scene.sensors["frame_transformer"].data.target_pos_w.torch.squeeze() + target_quat_w_tf = scene.sensors["frame_transformer"].data.target_quat_w.torch.squeeze() target_frame_names = scene.sensors["frame_transformer"].data.target_frame_names cube_center_idx = target_frame_names.index("CUBE_CENTER") @@ -525,7 +524,7 @@ def test_frame_transformer_all_bodies(sim): reordering_indices = [target_frame_names.index(name) for name in articulation_body_names] # default joint targets - default_actions = wp.to_torch(scene.articulations["robot"].data.default_joint_pos).clone() + default_actions = scene.articulations["robot"].data.default_joint_pos.torch.clone() # Define simulation stepping sim_dt = sim.get_physics_dt() # Simulate physics @@ -533,10 +532,10 @@ def test_frame_transformer_all_bodies(sim): # # reset if count % 25 == 0: # reset root state - root_state = wp.to_torch(scene.articulations["robot"].data.default_root_state).clone() + root_state = scene.articulations["robot"].data.default_root_state.torch.clone() root_state[:, :3] += scene.env_origins - joint_pos = scene.articulations["robot"].data.default_joint_pos - joint_vel = scene.articulations["robot"].data.default_joint_vel + joint_pos = scene.articulations["robot"].data.default_joint_pos.torch + joint_vel = scene.articulations["robot"].data.default_joint_vel.torch # -- set root state # -- robot scene.articulations["robot"].write_root_pose_to_sim(root_state[:, :7]) @@ -557,15 +556,15 @@ def test_frame_transformer_all_bodies(sim): # check absolute frame transforms in world frame # -- ground-truth - root_pose_w = wp.to_torch(scene.articulations["robot"].data.root_pose_w) - bodies_pos_w_gt = wp.to_torch(scene.articulations["robot"].data.body_pos_w) - bodies_quat_w_gt = wp.to_torch(scene.articulations["robot"].data.body_quat_w) + root_pose_w = scene.articulations["robot"].data.root_pose_w.torch + bodies_pos_w_gt = scene.articulations["robot"].data.body_pos_w.torch + bodies_quat_w_gt = scene.articulations["robot"].data.body_quat_w.torch # -- frame transformer - source_pos_w_tf = wp.to_torch(scene.sensors["frame_transformer"].data.source_pos_w) - source_quat_w_tf = wp.to_torch(scene.sensors["frame_transformer"].data.source_quat_w) - bodies_pos_w_tf = wp.to_torch(scene.sensors["frame_transformer"].data.target_pos_w) - bodies_quat_w_tf = wp.to_torch(scene.sensors["frame_transformer"].data.target_quat_w) + source_pos_w_tf = scene.sensors["frame_transformer"].data.source_pos_w.torch + source_quat_w_tf = scene.sensors["frame_transformer"].data.source_quat_w.torch + bodies_pos_w_tf = scene.sensors["frame_transformer"].data.target_pos_w.torch + bodies_quat_w_tf = scene.sensors["frame_transformer"].data.target_quat_w.torch # check if they are same torch.testing.assert_close(root_pose_w[:, :3], source_pos_w_tf) @@ -573,8 +572,8 @@ def test_frame_transformer_all_bodies(sim): torch.testing.assert_close(bodies_pos_w_gt, bodies_pos_w_tf[:, reordering_indices]) torch.testing.assert_close(bodies_quat_w_gt, bodies_quat_w_tf[:, reordering_indices]) - bodies_pos_source_tf = wp.to_torch(scene.sensors["frame_transformer"].data.target_pos_source) - bodies_quat_source_tf = wp.to_torch(scene.sensors["frame_transformer"].data.target_quat_source) + bodies_pos_source_tf = scene.sensors["frame_transformer"].data.target_pos_source.torch + bodies_quat_source_tf = scene.sensors["frame_transformer"].data.target_quat_source.torch # Go through each body and check if relative transforms are same for index in range(len(articulation_body_names)): @@ -728,22 +727,22 @@ class MultiRobotSceneCfg(InteractiveSceneCfg): # Reset periodically if count % 10 == 0: # Reset robot - root_state = wp.to_torch(scene.articulations["robot"].data.default_root_state).clone() + root_state = scene.articulations["robot"].data.default_root_state.torch.clone() root_state[:, :3] += scene.env_origins scene.articulations["robot"].write_root_pose_to_sim(root_state[:, :7]) scene.articulations["robot"].write_root_velocity_to_sim(root_state[:, 7:]) scene.articulations["robot"].write_joint_state_to_sim( - scene.articulations["robot"].data.default_joint_pos, - scene.articulations["robot"].data.default_joint_vel, + scene.articulations["robot"].data.default_joint_pos.torch, + scene.articulations["robot"].data.default_joint_vel.torch, ) # Reset robot_1 - root_state_1 = wp.to_torch(scene.articulations["robot_1"].data.default_root_state).clone() + root_state_1 = scene.articulations["robot_1"].data.default_root_state.torch.clone() root_state_1[:, :3] += scene.env_origins scene.articulations["robot_1"].write_root_pose_to_sim(root_state_1[:, :7]) scene.articulations["robot_1"].write_root_velocity_to_sim(root_state_1[:, 7:]) scene.articulations["robot_1"].write_joint_state_to_sim( - scene.articulations["robot_1"].data.default_joint_pos, - scene.articulations["robot_1"].data.default_joint_vel, + scene.articulations["robot_1"].data.default_joint_pos.torch, + scene.articulations["robot_1"].data.default_joint_vel.torch, ) scene.reset() @@ -756,21 +755,21 @@ class MultiRobotSceneCfg(InteractiveSceneCfg): # Get frame transformer data frame_transformer_data = scene.sensors["frame_transformer"].data - source_pos_w = wp.to_torch(frame_transformer_data.source_pos_w) - source_quat_w = wp.to_torch(frame_transformer_data.source_quat_w) - target_pos_w = wp.to_torch(frame_transformer_data.target_pos_w) + source_pos_w = frame_transformer_data.source_pos_w.torch + source_quat_w = frame_transformer_data.source_quat_w.torch + target_pos_w = frame_transformer_data.target_pos_w.torch # Get ground truth positions and orientations (after scene.update() so they're current) - robot_lf_pos_w = wp.to_torch(scene.articulations["robot"].data.body_pos_w)[:, robot_lf_shank_body_idx] - robot_1_lf_pos_w = wp.to_torch(scene.articulations["robot_1"].data.body_pos_w)[:, robot_1_lf_shank_body_idx] - robot_rf_pos_w = wp.to_torch(scene.articulations["robot"].data.body_pos_w)[:, robot_rf_shank_body_idx] - robot_1_rf_pos_w = wp.to_torch(scene.articulations["robot_1"].data.body_pos_w)[:, robot_1_rf_shank_body_idx] + robot_lf_pos_w = scene.articulations["robot"].data.body_pos_w.torch[:, robot_lf_shank_body_idx] + robot_1_lf_pos_w = scene.articulations["robot_1"].data.body_pos_w.torch[:, robot_1_lf_shank_body_idx] + robot_rf_pos_w = scene.articulations["robot"].data.body_pos_w.torch[:, robot_rf_shank_body_idx] + robot_1_rf_pos_w = scene.articulations["robot_1"].data.body_pos_w.torch[:, robot_1_rf_shank_body_idx] # Get expected source frame positions and orientations (after scene.update() so they're current) - expected_source_base_pos_w = wp.to_torch(scene.articulations[expected_source_robot].data.body_pos_w)[ + expected_source_base_pos_w = scene.articulations[expected_source_robot].data.body_pos_w.torch[ :, expected_source_base_body_idx ] - expected_source_base_quat_w = wp.to_torch(scene.articulations[expected_source_robot].data.body_quat_w)[ + expected_source_base_quat_w = scene.articulations[expected_source_robot].data.body_quat_w.torch[ :, expected_source_base_body_idx ] diff --git a/source/isaaclab_newton/test/sensors/test_imu.py b/source/isaaclab_newton/test/sensors/test_imu.py index 7bfd0ee4fba7..3edb4c908398 100644 --- a/source/isaaclab_newton/test/sensors/test_imu.py +++ b/source/isaaclab_newton/test/sensors/test_imu.py @@ -87,8 +87,8 @@ def test_data_shapes(sim): scene.update(sim.get_physics_dt()) imu: Imu = scene["imu"] - ang_vel = wp.to_torch(imu.data.ang_vel_b) - lin_acc = wp.to_torch(imu.data.lin_acc_b) + ang_vel = imu.data.ang_vel_b.torch + lin_acc = imu.data.lin_acc_b.torch assert ang_vel.shape == (2, 3) assert lin_acc.shape == (2, 3) @@ -106,7 +106,7 @@ def test_gravity_at_rest(sim): scene.update(sim.get_physics_dt()) imu: Imu = scene["imu"] - lin_acc = wp.to_torch(imu.data.lin_acc_b) + lin_acc = imu.data.lin_acc_b.torch # At rest, accelerometer should read ~9.81 in the up direction (Z body frame) torch.testing.assert_close( @@ -135,7 +135,7 @@ def test_angular_velocity_at_rest(sim): scene.update(sim.get_physics_dt()) imu: Imu = scene["imu"] - ang_vel = wp.to_torch(imu.data.ang_vel_b) + ang_vel = imu.data.ang_vel_b.torch torch.testing.assert_close( ang_vel, @@ -160,7 +160,7 @@ def test_reset(sim): imu: Imu = scene["imu"] - lin_acc = wp.to_torch(imu.data.lin_acc_b) + lin_acc = imu.data.lin_acc_b.torch assert torch.any(lin_acc != 0), "Expected non-zero data before reset" imu.reset() @@ -209,7 +209,7 @@ def test_freefall_acceleration(sim): scene.update(sim.get_physics_dt()) imu: Imu = scene["imu"] - lin_acc = wp.to_torch(imu.data.lin_acc_b) + lin_acc = imu.data.lin_acc_b.torch # In freefall, accelerometer should read near zero (gravity and inertial acceleration cancel) acc_magnitude = torch.norm(lin_acc, dim=-1) diff --git a/source/isaaclab_newton/test/sensors/test_pva.py b/source/isaaclab_newton/test/sensors/test_pva.py index c69a018f765b..8d14b9cb77ef 100644 --- a/source/isaaclab_newton/test/sensors/test_pva.py +++ b/source/isaaclab_newton/test/sensors/test_pva.py @@ -93,14 +93,14 @@ def test_data_shapes(sim): scene.update(sim.get_physics_dt()) pva: Pva = scene["pva"] - assert wp.to_torch(pva.data.pos_w).shape == (2, 3) - assert wp.to_torch(pva.data.quat_w).shape == (2, 4) - assert wp.to_torch(pva.data.pose_w).shape == (2, 7) - assert wp.to_torch(pva.data.lin_vel_b).shape == (2, 3) - assert wp.to_torch(pva.data.ang_vel_b).shape == (2, 3) - assert wp.to_torch(pva.data.lin_acc_b).shape == (2, 3) - assert wp.to_torch(pva.data.ang_acc_b).shape == (2, 3) - assert wp.to_torch(pva.data.projected_gravity_b).shape == (2, 3) + assert pva.data.pos_w.torch.shape == (2, 3) + assert pva.data.quat_w.torch.shape == (2, 4) + assert pva.data.pose_w.torch.shape == (2, 7) + assert pva.data.lin_vel_b.torch.shape == (2, 3) + assert pva.data.ang_vel_b.torch.shape == (2, 3) + assert pva.data.lin_acc_b.torch.shape == (2, 3) + assert pva.data.ang_acc_b.torch.shape == (2, 3) + assert pva.data.projected_gravity_b.torch.shape == (2, 3) def test_gravity_at_rest(sim): @@ -115,7 +115,7 @@ def test_gravity_at_rest(sim): scene.update(sim.get_physics_dt()) pva: Pva = scene["pva"] - proj_grav = wp.to_torch(pva.data.projected_gravity_b) + proj_grav = pva.data.projected_gravity_b.torch expected = torch.tensor([[0.0, 0.0, -1.0]], dtype=proj_grav.dtype, device=proj_grav.device).repeat(2, 1) torch.testing.assert_close(proj_grav, expected, atol=0.05, rtol=0.0) @@ -132,8 +132,8 @@ def test_velocity_at_rest(sim): scene.update(sim.get_physics_dt()) pva: Pva = scene["pva"] - lin_vel = wp.to_torch(pva.data.lin_vel_b) - ang_vel = wp.to_torch(pva.data.ang_vel_b) + lin_vel = pva.data.lin_vel_b.torch + ang_vel = pva.data.ang_vel_b.torch torch.testing.assert_close(lin_vel, torch.zeros_like(lin_vel), atol=0.05, rtol=0.0) torch.testing.assert_close(ang_vel, torch.zeros_like(ang_vel), atol=0.05, rtol=0.0) @@ -149,7 +149,7 @@ def test_position_nonzero(sim): scene.update(sim.get_physics_dt()) pva: Pva = scene["pva"] - pos = wp.to_torch(pva.data.pos_w) + pos = pva.data.pos_w.torch assert torch.all(pos[:, 2] > 0.0), f"Expected positive z position, got {pos[:, 2]}" @@ -166,7 +166,7 @@ def test_reset(sim): pva: Pva = scene["pva"] - pos = wp.to_torch(pva.data.pos_w) + pos = pva.data.pos_w.torch assert torch.any(pos != 0), "Expected non-zero data before reset" pva.reset() @@ -223,7 +223,7 @@ def test_freefall_velocity_increases(sim): scene.update(sim.get_physics_dt()) pva: Pva = scene["pva"] - lin_vel = wp.to_torch(pva.data.lin_vel_b) + lin_vel = pva.data.lin_vel_b.torch speed = torch.norm(lin_vel, dim=-1) assert torch.all(speed > 0.1), f"Expected non-zero velocity in freefall, got {speed}" @@ -245,8 +245,8 @@ def test_freefall_acceleration(sim): scene.update(sim.get_physics_dt()) pva: Pva = scene["pva"] - lin_acc = wp.to_torch(pva.data.lin_acc_b) - ang_acc = wp.to_torch(pva.data.ang_acc_b) + lin_acc = pva.data.lin_acc_b.torch + ang_acc = pva.data.ang_acc_b.torch # Coordinate acceleration in freefall should be ~(0, 0, -9.81) in body frame. expected_acc = torch.tensor([[0.0, 0.0, -9.81]], dtype=lin_acc.dtype, device=lin_acc.device).repeat(2, 1) @@ -315,8 +315,8 @@ def test_offset_and_rotated_body(sim): scene.update(sim.get_physics_dt()) pva: Pva = scene["pva"] - pos = wp.to_torch(pva.data.pos_w) - proj_grav = wp.to_torch(pva.data.projected_gravity_b) + pos = pva.data.pos_w.torch + proj_grav = pva.data.projected_gravity_b.torch # pos_w is in absolute world frame; subtract env origins to get env-relative position. # Expected env-relative: body at (0,0,5) + R_x(90) * (0,0,0.5) = (0, -0.5, 5) diff --git a/source/isaaclab_newton/test/sim/test_views_xform_prim_newton.py b/source/isaaclab_newton/test/sim/test_views_xform_prim_newton.py index 9785b6d62e2e..e8b28aad2d95 100644 --- a/source/isaaclab_newton/test/sim/test_views_xform_prim_newton.py +++ b/source/isaaclab_newton/test/sim/test_views_xform_prim_newton.py @@ -165,7 +165,7 @@ def test_world_attached_returns_initial_pose(device): sim.reset() view = FrameView("/World/StaticMarker", device=device) - pos = wp.to_torch(view.get_world_poses()[0]) + pos = view.get_world_poses()[0].torch expected = torch.tensor([list(WORLD_MARKER_POS)], device=device) torch.testing.assert_close(pos, expected, atol=1e-5, rtol=0) ctx.__exit__(None, None, None) @@ -193,6 +193,6 @@ def test_world_attached_set_world_roundtrip(device): view.set_world_poses(new_pos, new_quat) ret_pos, ret_quat = view.get_world_poses() - torch.testing.assert_close(wp.to_torch(ret_pos), wp.to_torch(new_pos), atol=1e-5, rtol=0) - torch.testing.assert_close(wp.to_torch(ret_quat), wp.to_torch(new_quat), atol=1e-5, rtol=0) + torch.testing.assert_close(ret_pos.torch, wp.to_torch(new_pos), atol=1e-5, rtol=0) + torch.testing.assert_close(ret_quat.torch, wp.to_torch(new_quat), atol=1e-5, rtol=0) ctx.__exit__(None, None, None) diff --git a/source/isaaclab_ovphysx/config/extension.toml b/source/isaaclab_ovphysx/config/extension.toml index ed4f5b39fb70..8648b7fa9587 100644 --- a/source/isaaclab_ovphysx/config/extension.toml +++ b/source/isaaclab_ovphysx/config/extension.toml @@ -1,7 +1,7 @@ [package] # Note: Semantic Versioning is used: https://semver.org/ -version = "0.1.1" +version = "0.1.2" # Description title = "OvPhysX simulation interfaces for IsaacLab core package" diff --git a/source/isaaclab_ovphysx/docs/CHANGELOG.rst b/source/isaaclab_ovphysx/docs/CHANGELOG.rst index 750a0397f23d..b2eb969d7845 100644 --- a/source/isaaclab_ovphysx/docs/CHANGELOG.rst +++ b/source/isaaclab_ovphysx/docs/CHANGELOG.rst @@ -1,6 +1,21 @@ Changelog --------- +0.1.2 (2026-04-23) +~~~~~~~~~~~~~~~~~~ + +Fixed +^^^^^ + +* Wrapped :attr:`~isaaclab_ovphysx.assets.ArticulationData.GRAVITY_VEC_W` and + :attr:`~isaaclab_ovphysx.assets.ArticulationData.FORWARD_VEC_B` in + :class:`~isaaclab.utils.warp.ProxyArray` to match the PhysX and Newton + backends. Public observations such as + :func:`~isaaclab.envs.mdp.observations.projected_gravity` access + ``asset.data.GRAVITY_VEC_W.torch``; the previous raw ``wp.array`` lacked + ``.torch`` and raised ``AttributeError`` on the ovphysx backend. + + 0.1.1 (2026-04-21) ~~~~~~~~~~~~~~~~~~~ diff --git a/source/isaaclab_ovphysx/isaaclab_ovphysx/assets/articulation/articulation.py b/source/isaaclab_ovphysx/isaaclab_ovphysx/assets/articulation/articulation.py index 7224d53d40ea..f3919d56da73 100644 --- a/source/isaaclab_ovphysx/isaaclab_ovphysx/assets/articulation/articulation.py +++ b/source/isaaclab_ovphysx/isaaclab_ovphysx/assets/articulation/articulation.py @@ -2006,9 +2006,9 @@ def _apply_actuator_model(self) -> None: all_joints = len(jids_t) == self._num_joints # warp -> torch (zero-copy on same device via DLPack) - jp_target_full = wp.to_torch(self._data.joint_pos_target) - jv_target_full = wp.to_torch(self._data.joint_vel_target) - je_target_full = wp.to_torch(self._data.joint_effort_target) + jp_target_full = self._data.joint_pos_target.torch + jv_target_full = self._data.joint_vel_target.torch + je_target_full = self._data.joint_effort_target.torch jp_target = jp_target_full if all_joints else jp_target_full[:, jids_t] jv_target = jv_target_full if all_joints else jv_target_full[:, jids_t] je_target = je_target_full if all_joints else je_target_full[:, jids_t] @@ -2019,8 +2019,8 @@ def _apply_actuator_model(self) -> None: joint_efforts=je_target, ) - jp_cur_full = wp.to_torch(self._data.joint_pos) - jv_cur_full = wp.to_torch(self._data.joint_vel) + jp_cur_full = self._data.joint_pos.torch + jv_cur_full = self._data.joint_vel.torch jp_cur = jp_cur_full if all_joints else jp_cur_full[:, jids_t] jv_cur = jv_cur_full if all_joints else jv_cur_full[:, jids_t] @@ -2057,14 +2057,14 @@ def format_limits(_, v: tuple[float, float]) -> str: return f"[{v[0]:.1e}, {v[1]:.1e}]" return f"[{v[0]:.3f}, {v[1]:.3f}]" - stiffnesses = self.data.joint_stiffness.numpy()[0].tolist() - dampings = self.data.joint_damping.numpy()[0].tolist() - armatures = self.data.joint_armature.numpy()[0].tolist() - frictions = self.data.joint_friction_coeff.numpy()[0].tolist() - pos_limits_np = self.data.joint_pos_limits.numpy().reshape(self._num_instances, self._num_joints, 2) + stiffnesses = self.data.joint_stiffness.warp.numpy()[0].tolist() + dampings = self.data.joint_damping.warp.numpy()[0].tolist() + armatures = self.data.joint_armature.warp.numpy()[0].tolist() + frictions = self.data.joint_friction_coeff.warp.numpy()[0].tolist() + pos_limits_np = self.data.joint_pos_limits.warp.numpy().reshape(self._num_instances, self._num_joints, 2) position_limits = [tuple(pos_limits_np[0, j].tolist()) for j in range(self._num_joints)] - velocity_limits = self.data.joint_vel_limits.numpy()[0].tolist() - effort_limits = self.data.joint_effort_limits.numpy()[0].tolist() + velocity_limits = self.data.joint_vel_limits.warp.numpy()[0].tolist() + effort_limits = self.data.joint_effort_limits.warp.numpy()[0].tolist() joint_table = PrettyTable() joint_table.title = f"Simulation Joint Information (Prim path: {self.cfg.prim_path})" @@ -2105,15 +2105,15 @@ def format_limits(_, v: tuple[float, float]) -> str: logger.info(f"Simulation parameters for joints in {self.cfg.prim_path}:\n" + joint_table.get_string()) if self.num_fixed_tendons > 0: - ft_stiffnesses = self.data.fixed_tendon_stiffness.numpy()[0].tolist() - ft_dampings = self.data.fixed_tendon_damping.numpy()[0].tolist() - ft_limit_stiffnesses = self.data.fixed_tendon_limit_stiffness.numpy()[0].tolist() - ft_limits_np = self.data.fixed_tendon_pos_limits.numpy().reshape( + ft_stiffnesses = self.data.fixed_tendon_stiffness.warp.numpy()[0].tolist() + ft_dampings = self.data.fixed_tendon_damping.warp.numpy()[0].tolist() + ft_limit_stiffnesses = self.data.fixed_tendon_limit_stiffness.warp.numpy()[0].tolist() + ft_limits_np = self.data.fixed_tendon_pos_limits.warp.numpy().reshape( self._num_instances, self.num_fixed_tendons, 2 ) ft_limits = [tuple(ft_limits_np[0, t].tolist()) for t in range(self.num_fixed_tendons)] - ft_rest_lengths = self.data.fixed_tendon_rest_length.numpy()[0].tolist() - ft_offsets = self.data.fixed_tendon_offset.numpy()[0].tolist() + ft_rest_lengths = self.data.fixed_tendon_rest_length.warp.numpy()[0].tolist() + ft_offsets = self.data.fixed_tendon_offset.warp.numpy()[0].tolist() tendon_table = PrettyTable() tendon_table.title = f"Simulation Fixed Tendon Information (Prim path: {self.cfg.prim_path})" @@ -2149,10 +2149,10 @@ def format_limits(_, v: tuple[float, float]) -> str: ) if self.num_spatial_tendons > 0: - st_stiffnesses = self.data.spatial_tendon_stiffness.numpy()[0].tolist() - st_dampings = self.data.spatial_tendon_damping.numpy()[0].tolist() - st_limit_stiffnesses = self.data.spatial_tendon_limit_stiffness.numpy()[0].tolist() - st_offsets = self.data.spatial_tendon_offset.numpy()[0].tolist() + st_stiffnesses = self.data.spatial_tendon_stiffness.warp.numpy()[0].tolist() + st_dampings = self.data.spatial_tendon_damping.warp.numpy()[0].tolist() + st_limit_stiffnesses = self.data.spatial_tendon_limit_stiffness.warp.numpy()[0].tolist() + st_offsets = self.data.spatial_tendon_offset.warp.numpy()[0].tolist() tendon_table = PrettyTable() tendon_table.title = f"Simulation Spatial Tendon Information (Prim path: {self.cfg.prim_path})" diff --git a/source/isaaclab_ovphysx/isaaclab_ovphysx/assets/articulation/articulation_data.py b/source/isaaclab_ovphysx/isaaclab_ovphysx/assets/articulation/articulation_data.py index 84c06675a022..fe264924c184 100644 --- a/source/isaaclab_ovphysx/isaaclab_ovphysx/assets/articulation/articulation_data.py +++ b/source/isaaclab_ovphysx/isaaclab_ovphysx/assets/articulation/articulation_data.py @@ -15,6 +15,7 @@ from isaaclab.assets.articulation.base_articulation_data import BaseArticulationData from isaaclab.utils.buffers import TimestampedBufferWarp as TimestampedBuffer +from isaaclab.utils.warp import ProxyArray from isaaclab_ovphysx import tensor_types as TT @@ -103,8 +104,8 @@ def __init__(self, bindings: dict[int, Any], device: str, binding_getter=None): gravity_dir_tiled = np.tile(gravity_dir, (self._num_instances, 1)) forward_tiled = np.tile(np.array([1.0, 0.0, 0.0], dtype=np.float32), (self._num_instances, 1)) - self.GRAVITY_VEC_W = wp.from_numpy(gravity_dir_tiled, dtype=wp.vec3f, device=device) - self.FORWARD_VEC_B = wp.from_numpy(forward_tiled, dtype=wp.vec3f, device=device) + self.GRAVITY_VEC_W = ProxyArray(wp.from_numpy(gravity_dir_tiled, dtype=wp.vec3f, device=device)) + self.FORWARD_VEC_B = ProxyArray(wp.from_numpy(forward_tiled, dtype=wp.vec3f, device=device)) def update(self, dt: float) -> None: """Update the data for the articulation. @@ -169,13 +170,15 @@ def is_primed(self, value: bool) -> None: """ @property - def default_root_pose(self) -> wp.array: + def default_root_pose(self) -> ProxyArray: """Default root pose ``[pos, quat]`` in the local environment frame. The position and quaternion are of the articulation root's actor frame. Shape is (num_instances,), dtype = wp.transformf. In torch this resolves to (num_instances, 7). """ - return self._default_root_pose + if self._default_root_pose_ta is None: + self._default_root_pose_ta = ProxyArray(self._default_root_pose) + return self._default_root_pose_ta @default_root_pose.setter def default_root_pose(self, value: wp.array) -> None: @@ -192,13 +195,15 @@ def default_root_pose(self, value: wp.array) -> None: self._default_root_pose.assign(value) @property - def default_root_vel(self) -> wp.array: + def default_root_vel(self) -> ProxyArray: """Default root velocity ``[lin_vel, ang_vel]`` in the local environment frame. The linear and angular velocities are of the articulation root's center of mass frame. Shape is (num_instances,), dtype = wp.spatial_vectorf. In torch this resolves to (num_instances, 6). """ - return self._default_root_vel + if self._default_root_vel_ta is None: + self._default_root_vel_ta = ProxyArray(self._default_root_vel) + return self._default_root_vel_ta @default_root_vel.setter def default_root_vel(self, value: wp.array) -> None: @@ -215,7 +220,7 @@ def default_root_vel(self, value: wp.array) -> None: self._default_root_vel.assign(value) @property - def default_joint_pos(self) -> wp.array: + def default_joint_pos(self) -> ProxyArray: """Default joint positions of all joints [m or rad, depending on joint type]. Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to @@ -223,7 +228,9 @@ def default_joint_pos(self) -> wp.array: This quantity is configured through the :attr:`isaaclab.assets.ArticulationCfg.init_state` parameter. """ - return self._default_joint_pos + if self._default_joint_pos_ta is None: + self._default_joint_pos_ta = ProxyArray(self._default_joint_pos) + return self._default_joint_pos_ta @default_joint_pos.setter def default_joint_pos(self, value: wp.array) -> None: @@ -240,7 +247,7 @@ def default_joint_pos(self, value: wp.array) -> None: self._default_joint_pos.assign(value) @property - def default_joint_vel(self) -> wp.array: + def default_joint_vel(self) -> ProxyArray: """Default joint velocities of all joints [m/s or rad/s, depending on joint type]. Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to @@ -248,7 +255,9 @@ def default_joint_vel(self) -> wp.array: This quantity is configured through the :attr:`isaaclab.assets.ArticulationCfg.init_state` parameter. """ - return self._default_joint_vel + if self._default_joint_vel_ta is None: + self._default_joint_vel_ta = ProxyArray(self._default_joint_vel) + return self._default_joint_vel_ta @default_joint_vel.setter def default_joint_vel(self, value: wp.array) -> None: @@ -269,7 +278,7 @@ def default_joint_vel(self, value: wp.array) -> None: """ @property - def joint_pos_target(self) -> wp.array: + def joint_pos_target(self) -> ProxyArray: """Joint position targets commanded by the user [m or rad, depending on joint type]. Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to @@ -279,10 +288,12 @@ def joint_pos_target(self) -> wp.array: For an explicit actuator model, the targets are used to compute the joint torques (see :attr:`applied_torque`), which are then set into the simulation. """ - return self._joint_pos_target + if self._joint_pos_target_ta is None: + self._joint_pos_target_ta = ProxyArray(self._joint_pos_target) + return self._joint_pos_target_ta @property - def joint_vel_target(self) -> wp.array: + def joint_vel_target(self) -> ProxyArray: """Joint velocity targets commanded by the user [m/s or rad/s, depending on joint type]. Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to @@ -292,10 +303,12 @@ def joint_vel_target(self) -> wp.array: For an explicit actuator model, the targets are used to compute the joint torques (see :attr:`applied_torque`), which are then set into the simulation. """ - return self._joint_vel_target + if self._joint_vel_target_ta is None: + self._joint_vel_target_ta = ProxyArray(self._joint_vel_target) + return self._joint_vel_target_ta @property - def joint_effort_target(self) -> wp.array: + def joint_effort_target(self) -> ProxyArray: """Joint effort targets commanded by the user [N or N*m, depending on joint type]. Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to @@ -305,14 +318,16 @@ def joint_effort_target(self) -> wp.array: For an explicit actuator model, the targets are used to compute the joint torques (see :attr:`applied_torque`), which are then set into the simulation. """ - return self._joint_effort_target + if self._joint_effort_target_ta is None: + self._joint_effort_target_ta = ProxyArray(self._joint_effort_target) + return self._joint_effort_target_ta """ Joint commands -- Explicit actuators. """ @property - def computed_torque(self) -> wp.array: + def computed_torque(self) -> ProxyArray: """Joint torques computed from the actuator model (before clipping) [N*m]. Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to @@ -322,10 +337,12 @@ def computed_torque(self) -> wp.array: It is exposed for users who want to inspect the computations inside the actuator model. For instance, to penalize the learning agent for a difference between the computed and applied torques. """ - return self._computed_torque + if self._computed_torque_ta is None: + self._computed_torque_ta = ProxyArray(self._computed_torque) + return self._computed_torque_ta @property - def applied_torque(self) -> wp.array: + def applied_torque(self) -> ProxyArray: """Joint torques applied from the actuator model (after clipping) [N*m]. Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to @@ -334,14 +351,16 @@ def applied_torque(self) -> wp.array: These torques are set into the simulation, after clipping the :attr:`computed_torque` based on the actuator model. """ - return self._applied_torque + if self._applied_torque_ta is None: + self._applied_torque_ta = ProxyArray(self._applied_torque) + return self._applied_torque_ta """ Joint properties """ @property - def joint_stiffness(self) -> wp.array: + def joint_stiffness(self) -> ProxyArray: """Joint stiffness provided to the simulation. Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to @@ -349,10 +368,12 @@ def joint_stiffness(self) -> wp.array: In the case of explicit actuators, the value for the corresponding joints is zero. """ - return self._joint_stiffness + if self._joint_stiffness_ta is None: + self._joint_stiffness_ta = ProxyArray(self._joint_stiffness) + return self._joint_stiffness_ta @property - def joint_damping(self) -> wp.array: + def joint_damping(self) -> ProxyArray: """Joint damping provided to the simulation. Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to @@ -360,28 +381,34 @@ def joint_damping(self) -> wp.array: In the case of explicit actuators, the value for the corresponding joints is zero. """ - return self._joint_damping + if self._joint_damping_ta is None: + self._joint_damping_ta = ProxyArray(self._joint_damping) + return self._joint_damping_ta @property - def joint_armature(self) -> wp.array: + def joint_armature(self) -> ProxyArray: """Joint armature provided to the simulation. Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to (num_instances, num_joints). """ - return self._joint_armature + if self._joint_armature_ta is None: + self._joint_armature_ta = ProxyArray(self._joint_armature) + return self._joint_armature_ta @property - def joint_friction_coeff(self) -> wp.array: + def joint_friction_coeff(self) -> ProxyArray: """Joint static friction coefficient provided to the simulation. Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to (num_instances, num_joints). """ - return self._joint_friction_coeff + if self._joint_friction_coeff_ta is None: + self._joint_friction_coeff_ta = ProxyArray(self._joint_friction_coeff) + return self._joint_friction_coeff_ta @property - def joint_pos_limits(self) -> wp.array: + def joint_pos_limits(self) -> ProxyArray: """Joint position limits provided to the simulation. Shape is (num_instances, num_joints), dtype = wp.vec2f. In torch this resolves to @@ -389,32 +416,38 @@ def joint_pos_limits(self) -> wp.array: The limits are in the order :math:`[lower, upper]`. """ - return self._joint_pos_limits + if self._joint_pos_limits_ta is None: + self._joint_pos_limits_ta = ProxyArray(self._joint_pos_limits) + return self._joint_pos_limits_ta @property - def joint_vel_limits(self) -> wp.array: + def joint_vel_limits(self) -> ProxyArray: """Joint maximum velocity provided to the simulation [m/s or rad/s, depending on joint type]. Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to (num_instances, num_joints). """ - return self._joint_vel_limits + if self._joint_vel_limits_ta is None: + self._joint_vel_limits_ta = ProxyArray(self._joint_vel_limits) + return self._joint_vel_limits_ta @property - def joint_effort_limits(self) -> wp.array: + def joint_effort_limits(self) -> ProxyArray: """Joint maximum effort provided to the simulation [N or N*m, depending on joint type]. Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to (num_instances, num_joints). """ - return self._joint_effort_limits + if self._joint_effort_limits_ta is None: + self._joint_effort_limits_ta = ProxyArray(self._joint_effort_limits) + return self._joint_effort_limits_ta """ Joint properties - Custom. """ @property - def soft_joint_pos_limits(self) -> wp.array: + def soft_joint_pos_limits(self) -> ProxyArray: r"""Soft joint position limits for all joints. Shape is (num_instances, num_joints), dtype = wp.vec2f. In torch this resolves to @@ -435,10 +468,12 @@ def soft_joint_pos_limits(self) -> wp.array: The soft joint position limits help specify a safety region around the joint limits. It isn't used by the simulation, but is useful for learning agents to prevent the joint positions from violating the limits. """ - return self._soft_joint_pos_limits + if self._soft_joint_pos_limits_ta is None: + self._soft_joint_pos_limits_ta = ProxyArray(self._soft_joint_pos_limits) + return self._soft_joint_pos_limits_ta @property - def soft_joint_vel_limits(self) -> wp.array: + def soft_joint_vel_limits(self) -> ProxyArray: """Soft joint velocity limits for all joints. Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to @@ -447,121 +482,145 @@ def soft_joint_vel_limits(self) -> wp.array: These are obtained from the actuator model. It may differ from :attr:`joint_vel_limits` if the actuator model has a variable velocity limit model. For instance, in a variable gear ratio actuator model. """ - return self._soft_joint_vel_limits + if self._soft_joint_vel_limits_ta is None: + self._soft_joint_vel_limits_ta = ProxyArray(self._soft_joint_vel_limits) + return self._soft_joint_vel_limits_ta @property - def gear_ratio(self) -> wp.array: + def gear_ratio(self) -> ProxyArray: """Gear ratio for relating motor torques to applied joint torques. Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to (num_instances, num_joints). """ - return self._gear_ratio + if self._gear_ratio_ta is None: + self._gear_ratio_ta = ProxyArray(self._gear_ratio) + return self._gear_ratio_ta """ Fixed tendon properties. """ @property - def fixed_tendon_stiffness(self) -> wp.array: + def fixed_tendon_stiffness(self) -> ProxyArray: """Fixed tendon stiffness provided to the simulation. Shape is (num_instances, num_fixed_tendons), dtype = wp.float32. In torch this resolves to (num_instances, num_fixed_tendons). """ - return self._fixed_tendon_stiffness + if self._fixed_tendon_stiffness_ta is None: + self._fixed_tendon_stiffness_ta = ProxyArray(self._fixed_tendon_stiffness) + return self._fixed_tendon_stiffness_ta @property - def fixed_tendon_damping(self) -> wp.array: + def fixed_tendon_damping(self) -> ProxyArray: """Fixed tendon damping provided to the simulation. Shape is (num_instances, num_fixed_tendons), dtype = wp.float32. In torch this resolves to (num_instances, num_fixed_tendons). """ - return self._fixed_tendon_damping + if self._fixed_tendon_damping_ta is None: + self._fixed_tendon_damping_ta = ProxyArray(self._fixed_tendon_damping) + return self._fixed_tendon_damping_ta @property - def fixed_tendon_limit_stiffness(self) -> wp.array: + def fixed_tendon_limit_stiffness(self) -> ProxyArray: """Fixed tendon limit stiffness provided to the simulation. Shape is (num_instances, num_fixed_tendons), dtype = wp.float32. In torch this resolves to (num_instances, num_fixed_tendons). """ - return self._fixed_tendon_limit_stiffness + if self._fixed_tendon_limit_stiffness_ta is None: + self._fixed_tendon_limit_stiffness_ta = ProxyArray(self._fixed_tendon_limit_stiffness) + return self._fixed_tendon_limit_stiffness_ta @property - def fixed_tendon_rest_length(self) -> wp.array: + def fixed_tendon_rest_length(self) -> ProxyArray: """Fixed tendon rest length provided to the simulation. Shape is (num_instances, num_fixed_tendons), dtype = wp.float32. In torch this resolves to (num_instances, num_fixed_tendons). """ - return self._fixed_tendon_rest_length + if self._fixed_tendon_rest_length_ta is None: + self._fixed_tendon_rest_length_ta = ProxyArray(self._fixed_tendon_rest_length) + return self._fixed_tendon_rest_length_ta @property - def fixed_tendon_offset(self) -> wp.array: + def fixed_tendon_offset(self) -> ProxyArray: """Fixed tendon offset provided to the simulation. Shape is (num_instances, num_fixed_tendons), dtype = wp.float32. In torch this resolves to (num_instances, num_fixed_tendons). """ - return self._fixed_tendon_offset + if self._fixed_tendon_offset_ta is None: + self._fixed_tendon_offset_ta = ProxyArray(self._fixed_tendon_offset) + return self._fixed_tendon_offset_ta @property - def fixed_tendon_pos_limits(self) -> wp.array: + def fixed_tendon_pos_limits(self) -> ProxyArray: """Fixed tendon position limits provided to the simulation. Shape is (num_instances, num_fixed_tendons), dtype = wp.vec2f. In torch this resolves to (num_instances, num_fixed_tendons, 2). """ - return self._fixed_tendon_pos_limits + if self._fixed_tendon_pos_limits_ta is None: + self._fixed_tendon_pos_limits_ta = ProxyArray(self._fixed_tendon_pos_limits) + return self._fixed_tendon_pos_limits_ta """ Spatial tendon properties. """ @property - def spatial_tendon_stiffness(self) -> wp.array: + def spatial_tendon_stiffness(self) -> ProxyArray: """Spatial tendon stiffness provided to the simulation. Shape is (num_instances, num_spatial_tendons), dtype = wp.float32. In torch this resolves to (num_instances, num_spatial_tendons). """ - return self._spatial_tendon_stiffness + if self._spatial_tendon_stiffness_ta is None: + self._spatial_tendon_stiffness_ta = ProxyArray(self._spatial_tendon_stiffness) + return self._spatial_tendon_stiffness_ta @property - def spatial_tendon_damping(self) -> wp.array: + def spatial_tendon_damping(self) -> ProxyArray: """Spatial tendon damping provided to the simulation. Shape is (num_instances, num_spatial_tendons), dtype = wp.float32. In torch this resolves to (num_instances, num_spatial_tendons). """ - return self._spatial_tendon_damping + if self._spatial_tendon_damping_ta is None: + self._spatial_tendon_damping_ta = ProxyArray(self._spatial_tendon_damping) + return self._spatial_tendon_damping_ta @property - def spatial_tendon_limit_stiffness(self) -> wp.array: + def spatial_tendon_limit_stiffness(self) -> ProxyArray: """Spatial tendon limit stiffness provided to the simulation. Shape is (num_instances, num_spatial_tendons), dtype = wp.float32. In torch this resolves to (num_instances, num_spatial_tendons). """ - return self._spatial_tendon_limit_stiffness + if self._spatial_tendon_limit_stiffness_ta is None: + self._spatial_tendon_limit_stiffness_ta = ProxyArray(self._spatial_tendon_limit_stiffness) + return self._spatial_tendon_limit_stiffness_ta @property - def spatial_tendon_offset(self) -> wp.array: + def spatial_tendon_offset(self) -> ProxyArray: """Spatial tendon offset provided to the simulation. Shape is (num_instances, num_spatial_tendons), dtype = wp.float32. In torch this resolves to (num_instances, num_spatial_tendons). """ - return self._spatial_tendon_offset + if self._spatial_tendon_offset_ta is None: + self._spatial_tendon_offset_ta = ProxyArray(self._spatial_tendon_offset) + return self._spatial_tendon_offset_ta """ Root state properties. """ @property - def root_link_pose_w(self) -> wp.array: + def root_link_pose_w(self) -> ProxyArray: """Root link pose ``[pos, quat]`` in simulation world frame. Shape is (num_instances,), dtype = wp.transformf. In torch this resolves to (num_instances, 7). @@ -569,10 +628,12 @@ def root_link_pose_w(self) -> wp.array: The orientation is provided in (x, y, z, w) format. """ self._read_transform_binding(TT.ROOT_POSE, self._root_link_pose_w) - return self._root_link_pose_w.data + if self._root_link_pose_w_ta is None: + self._root_link_pose_w_ta = ProxyArray(self._root_link_pose_w.data) + return self._root_link_pose_w_ta @property - def root_link_vel_w(self) -> wp.array: + def root_link_vel_w(self) -> ProxyArray: """Root link velocity ``[lin_vel, ang_vel]`` in simulation world frame. Shape is (num_instances,), dtype = wp.spatial_vectorf. In torch this resolves to (num_instances, 6). @@ -591,10 +652,12 @@ def root_link_vel_w(self) -> wp.array: device=self.device, ) self._root_link_vel_w.timestamp = self._sim_timestamp - return self._root_link_vel_w.data + if self._root_link_vel_w_ta is None: + self._root_link_vel_w_ta = ProxyArray(self._root_link_vel_w.data) + return self._root_link_vel_w_ta @property - def root_com_pose_w(self) -> wp.array: + def root_com_pose_w(self) -> ProxyArray: """Root center of mass pose ``[pos, quat]`` in simulation world frame. Shape is (num_instances,), dtype = wp.transformf. In torch this resolves to (num_instances, 7). @@ -610,10 +673,12 @@ def root_com_pose_w(self) -> wp.array: device=self.device, ) self._root_com_pose_w.timestamp = self._sim_timestamp - return self._root_com_pose_w.data + if self._root_com_pose_w_ta is None: + self._root_com_pose_w_ta = ProxyArray(self._root_com_pose_w.data) + return self._root_com_pose_w_ta @property - def root_com_vel_w(self) -> wp.array: + def root_com_vel_w(self) -> ProxyArray: """Root center of mass velocity ``[lin_vel, ang_vel]`` in simulation world frame. Shape is (num_instances,), dtype = wp.spatial_vectorf. In torch this resolves to (num_instances, 6). @@ -621,23 +686,27 @@ def root_com_vel_w(self) -> wp.array: relative to the world. """ self._read_spatial_vector_binding(TT.ROOT_VELOCITY, self._root_com_vel_w) - return self._root_com_vel_w.data + if self._root_com_vel_w_ta is None: + self._root_com_vel_w_ta = ProxyArray(self._root_com_vel_w.data) + return self._root_com_vel_w_ta """ Body state properties. """ @property - def body_mass(self) -> wp.array: + def body_mass(self) -> ProxyArray: """Body mass in the world frame [kg]. Shape is (num_instances, num_bodies), dtype = wp.float32. In torch this resolves to (num_instances, num_bodies). """ - return self._body_mass + if self._body_mass_ta is None: + self._body_mass_ta = ProxyArray(self._body_mass) + return self._body_mass_ta @property - def body_inertia(self) -> wp.array: + def body_inertia(self) -> ProxyArray: """Flattened body inertia in the world frame [kg*m^2]. Shape is (num_instances, num_bodies, 9), dtype = wp.float32. In torch this resolves to @@ -645,10 +714,12 @@ def body_inertia(self) -> wp.array: Stored as a flattened 3x3 inertia matrix per body. """ - return self._body_inertia + if self._body_inertia_ta is None: + self._body_inertia_ta = ProxyArray(self._body_inertia) + return self._body_inertia_ta @property - def body_link_pose_w(self) -> wp.array: + def body_link_pose_w(self) -> ProxyArray: """Body link pose ``[pos, quat]`` in simulation world frame. Shape is (num_instances, num_bodies), dtype = wp.transformf. In torch this resolves to (num_instances, num_bodies, 7). @@ -657,10 +728,12 @@ def body_link_pose_w(self) -> wp.array: The orientation is provided in (x, y, z, w) format. """ self._read_transform_binding(TT.LINK_POSE, self._body_link_pose_w) - return self._body_link_pose_w.data + if self._body_link_pose_w_ta is None: + self._body_link_pose_w_ta = ProxyArray(self._body_link_pose_w.data) + return self._body_link_pose_w_ta @property - def body_link_vel_w(self) -> wp.array: + def body_link_vel_w(self) -> ProxyArray: """Body link velocity ``[lin_vel, ang_vel]`` in simulation world frame. Shape is (num_instances, num_bodies), dtype = wp.spatial_vectorf. In torch this resolves to (num_instances, num_bodies, 6). @@ -669,10 +742,12 @@ def body_link_vel_w(self) -> wp.array: relative to the world. """ self._read_spatial_vector_binding(TT.LINK_VELOCITY, self._body_link_vel_w) - return self._body_link_vel_w.data + if self._body_link_vel_w_ta is None: + self._body_link_vel_w_ta = ProxyArray(self._body_link_vel_w.data) + return self._body_link_vel_w_ta @property - def body_com_pose_w(self) -> wp.array: + def body_com_pose_w(self) -> ProxyArray: """Body center of mass pose ``[pos, quat]`` in simulation world frame. Shape is (num_instances, num_bodies), dtype = wp.transformf. In torch this resolves to (num_instances, num_bodies, 7). @@ -689,10 +764,12 @@ def body_com_pose_w(self) -> wp.array: device=self.device, ) self._body_com_pose_w.timestamp = self._sim_timestamp - return self._body_com_pose_w.data + if self._body_com_pose_w_ta is None: + self._body_com_pose_w_ta = ProxyArray(self._body_com_pose_w.data) + return self._body_com_pose_w_ta @property - def body_com_vel_w(self) -> wp.array: + def body_com_vel_w(self) -> ProxyArray: """Body center of mass velocity ``[lin_vel, ang_vel]`` in simulation world frame. Shape is (num_instances, num_bodies), dtype = wp.spatial_vectorf. In torch this resolves to (num_instances, num_bodies, 6). @@ -707,7 +784,7 @@ def body_com_vel_w(self) -> wp.array: return self.body_link_vel_w @property - def body_com_acc_w(self) -> wp.array: + def body_com_acc_w(self) -> ProxyArray: """Acceleration of all bodies center of mass ``[lin_acc, ang_acc]`` [m/s^2, rad/s^2]. Shape is (num_instances, num_bodies), dtype = wp.spatial_vectorf. In torch this resolves to (num_instances, num_bodies, 6). @@ -715,10 +792,12 @@ def body_com_acc_w(self) -> wp.array: All values are relative to the world. """ self._read_spatial_vector_binding(TT.LINK_ACCELERATION, self._body_com_acc_w) - return self._body_com_acc_w.data + if self._body_com_acc_w_ta is None: + self._body_com_acc_w_ta = ProxyArray(self._body_com_acc_w.data) + return self._body_com_acc_w_ta @property - def body_com_pose_b(self) -> wp.array: + def body_com_pose_b(self) -> ProxyArray: """Center of mass pose ``[pos, quat]`` of all bodies in their respective body's link frames. Shape is (num_instances, num_bodies), dtype = wp.transformf. In torch this resolves to (num_instances, num_bodies, 7). @@ -727,10 +806,12 @@ def body_com_pose_b(self) -> wp.array: The orientation is provided in (x, y, z, w) format. """ self._read_transform_binding(TT.BODY_COM_POSE, self._body_com_pose_b) - return self._body_com_pose_b.data + if self._body_com_pose_b_ta is None: + self._body_com_pose_b_ta = ProxyArray(self._body_com_pose_b.data) + return self._body_com_pose_b_ta @property - def body_incoming_joint_wrench_b(self) -> wp.array: + def body_incoming_joint_wrench_b(self) -> ProxyArray: """Incoming joint wrenches on each body in the body frame [N, N*m]. Shape is (num_instances, num_bodies), dtype = wp.spatial_vectorf. In torch this resolves to @@ -742,34 +823,40 @@ def body_incoming_joint_wrench_b(self) -> wp.array: TT.LINK_INCOMING_JOINT_FORCE, self._body_incoming_joint_wrench_buf, ) - return self._body_incoming_joint_wrench_buf.data + if self._body_incoming_joint_wrench_b_ta is None: + self._body_incoming_joint_wrench_b_ta = ProxyArray(self._body_incoming_joint_wrench_buf.data) + return self._body_incoming_joint_wrench_b_ta """ Joint state properties. """ @property - def joint_pos(self) -> wp.array: + def joint_pos(self) -> ProxyArray: """Joint positions of all joints [m or rad, depending on joint type]. Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to (num_instances, num_joints). """ self._read_binding_into_buf(TT.DOF_POSITION, self._joint_pos_buf) - return self._joint_pos_buf.data + if self._joint_pos_ta is None: + self._joint_pos_ta = ProxyArray(self._joint_pos_buf.data) + return self._joint_pos_ta @property - def joint_vel(self) -> wp.array: + def joint_vel(self) -> ProxyArray: """Joint velocities of all joints [m/s or rad/s, depending on joint type]. Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to (num_instances, num_joints). """ self._read_binding_into_buf(TT.DOF_VELOCITY, self._joint_vel_buf) - return self._joint_vel_buf.data + if self._joint_vel_ta is None: + self._joint_vel_ta = ProxyArray(self._joint_vel_buf.data) + return self._joint_vel_ta @property - def joint_acc(self) -> wp.array: + def joint_acc(self) -> ProxyArray: """Joint acceleration of all joints [m/s^2 or rad/s^2, depending on joint type]. Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to @@ -778,14 +865,16 @@ def joint_acc(self) -> wp.array: .. note:: This quantity is computed via finite differencing of joint velocities. """ - return self._joint_acc.data + if self._joint_acc_ta is None: + self._joint_acc_ta = ProxyArray(self._joint_acc.data) + return self._joint_acc_ta """ Derived Properties. """ @property - def projected_gravity_b(self) -> wp.array: + def projected_gravity_b(self) -> ProxyArray: """Projection of the gravity direction on base frame. Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). """ @@ -798,10 +887,12 @@ def projected_gravity_b(self) -> wp.array: device=self.device, ) self._projected_gravity_b.timestamp = self._sim_timestamp - return self._projected_gravity_b.data + if self._projected_gravity_b_ta is None: + self._projected_gravity_b_ta = ProxyArray(self._projected_gravity_b.data) + return self._projected_gravity_b_ta @property - def heading_w(self) -> wp.array: + def heading_w(self) -> ProxyArray: """Yaw heading of the base frame (in radians). Shape is (num_instances,), dtype = wp.float32. @@ -818,10 +909,12 @@ def heading_w(self) -> wp.array: device=self.device, ) self._heading_w.timestamp = self._sim_timestamp - return self._heading_w.data + if self._heading_w_ta is None: + self._heading_w_ta = ProxyArray(self._heading_w.data) + return self._heading_w_ta @property - def root_link_lin_vel_b(self) -> wp.array: + def root_link_lin_vel_b(self) -> ProxyArray: """Root link linear velocity in base frame [m/s]. Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). @@ -836,10 +929,12 @@ def root_link_lin_vel_b(self) -> wp.array: device=self.device, ) self._root_link_lin_vel_b.timestamp = self._sim_timestamp - return self._root_link_lin_vel_b.data + if self._root_link_lin_vel_b_ta is None: + self._root_link_lin_vel_b_ta = ProxyArray(self._root_link_lin_vel_b.data) + return self._root_link_lin_vel_b_ta @property - def root_link_ang_vel_b(self) -> wp.array: + def root_link_ang_vel_b(self) -> ProxyArray: """Root link angular velocity in base frame [rad/s]. Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). @@ -854,10 +949,12 @@ def root_link_ang_vel_b(self) -> wp.array: device=self.device, ) self._root_link_ang_vel_b.timestamp = self._sim_timestamp - return self._root_link_ang_vel_b.data + if self._root_link_ang_vel_b_ta is None: + self._root_link_ang_vel_b_ta = ProxyArray(self._root_link_ang_vel_b.data) + return self._root_link_ang_vel_b_ta @property - def root_com_lin_vel_b(self) -> wp.array: + def root_com_lin_vel_b(self) -> ProxyArray: """Root center of mass linear velocity in base frame [m/s]. Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). @@ -873,10 +970,12 @@ def root_com_lin_vel_b(self) -> wp.array: device=self.device, ) self._root_com_lin_vel_b.timestamp = self._sim_timestamp - return self._root_com_lin_vel_b.data + if self._root_com_lin_vel_b_ta is None: + self._root_com_lin_vel_b_ta = ProxyArray(self._root_com_lin_vel_b.data) + return self._root_com_lin_vel_b_ta @property - def root_com_ang_vel_b(self) -> wp.array: + def root_com_ang_vel_b(self) -> ProxyArray: """Root center of mass angular velocity in base frame [rad/s]. Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). @@ -892,196 +991,255 @@ def root_com_ang_vel_b(self) -> wp.array: device=self.device, ) self._root_com_ang_vel_b.timestamp = self._sim_timestamp - return self._root_com_ang_vel_b.data + if self._root_com_ang_vel_b_ta is None: + self._root_com_ang_vel_b_ta = ProxyArray(self._root_com_ang_vel_b.data) + return self._root_com_ang_vel_b_ta """ Sliced properties. """ @property - def root_link_pos_w(self) -> wp.array: + def root_link_pos_w(self) -> ProxyArray: """Root link position in simulation world frame. Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). This quantity is the position of the actor frame of the root rigid body relative to the world. """ - return self._get_pos_from_transform(self.root_link_pose_w) + parent = self.root_link_pose_w + if self._root_link_pos_w_ta is None: + self._root_link_pos_w_ta = ProxyArray(self._get_pos_from_transform(parent.warp)) + return self._root_link_pos_w_ta @property - def root_link_quat_w(self) -> wp.array: + def root_link_quat_w(self) -> ProxyArray: """Root link orientation (x, y, z, w) in simulation world frame. Shape is (num_instances,), dtype = wp.quatf. In torch this resolves to (num_instances, 4). This quantity is the orientation of the actor frame of the root rigid body. """ - return self._get_quat_from_transform(self.root_link_pose_w) + parent = self.root_link_pose_w + if self._root_link_quat_w_ta is None: + self._root_link_quat_w_ta = ProxyArray(self._get_quat_from_transform(parent.warp)) + return self._root_link_quat_w_ta @property - def root_link_lin_vel_w(self) -> wp.array: + def root_link_lin_vel_w(self) -> ProxyArray: """Root linear velocity in simulation world frame [m/s]. Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). This quantity is the linear velocity of the root rigid body's actor frame relative to the world. """ - return self._get_lin_vel_from_spatial_vector(self.root_link_vel_w) + parent = self.root_link_vel_w + if self._root_link_lin_vel_w_ta is None: + self._root_link_lin_vel_w_ta = ProxyArray(self._get_lin_vel_from_spatial_vector(parent.warp)) + return self._root_link_lin_vel_w_ta @property - def root_link_ang_vel_w(self) -> wp.array: + def root_link_ang_vel_w(self) -> ProxyArray: """Root link angular velocity in simulation world frame [rad/s]. Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). This quantity is the angular velocity of the actor frame of the root rigid body relative to the world. """ - return self._get_ang_vel_from_spatial_vector(self.root_link_vel_w) + parent = self.root_link_vel_w + if self._root_link_ang_vel_w_ta is None: + self._root_link_ang_vel_w_ta = ProxyArray(self._get_ang_vel_from_spatial_vector(parent.warp)) + return self._root_link_ang_vel_w_ta @property - def root_com_pos_w(self) -> wp.array: + def root_com_pos_w(self) -> ProxyArray: """Root center of mass position in simulation world frame. Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). This quantity is the position of the center of mass frame of the root rigid body relative to the world. """ - return self._get_pos_from_transform(self.root_com_pose_w) + parent = self.root_com_pose_w + if self._root_com_pos_w_ta is None: + self._root_com_pos_w_ta = ProxyArray(self._get_pos_from_transform(parent.warp)) + return self._root_com_pos_w_ta @property - def root_com_quat_w(self) -> wp.array: + def root_com_quat_w(self) -> ProxyArray: """Root center of mass orientation (x, y, z, w) in simulation world frame. Shape is (num_instances,), dtype = wp.quatf. In torch this resolves to (num_instances, 4). This quantity is the orientation of the principal axes of inertia of the root rigid body relative to the world. """ - return self._get_quat_from_transform(self.root_com_pose_w) + parent = self.root_com_pose_w + if self._root_com_quat_w_ta is None: + self._root_com_quat_w_ta = ProxyArray(self._get_quat_from_transform(parent.warp)) + return self._root_com_quat_w_ta @property - def root_com_lin_vel_w(self) -> wp.array: + def root_com_lin_vel_w(self) -> ProxyArray: """Root center of mass linear velocity in simulation world frame [m/s]. Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). This quantity is the linear velocity of the root rigid body's center of mass frame relative to the world. """ - return self._get_lin_vel_from_spatial_vector(self.root_com_vel_w) + parent = self.root_com_vel_w + if self._root_com_lin_vel_w_ta is None: + self._root_com_lin_vel_w_ta = ProxyArray(self._get_lin_vel_from_spatial_vector(parent.warp)) + return self._root_com_lin_vel_w_ta @property - def root_com_ang_vel_w(self) -> wp.array: + def root_com_ang_vel_w(self) -> ProxyArray: """Root center of mass angular velocity in simulation world frame [rad/s]. Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). This quantity is the angular velocity of the root rigid body's center of mass frame relative to the world. """ - return self._get_ang_vel_from_spatial_vector(self.root_com_vel_w) + parent = self.root_com_vel_w + if self._root_com_ang_vel_w_ta is None: + self._root_com_ang_vel_w_ta = ProxyArray(self._get_ang_vel_from_spatial_vector(parent.warp)) + return self._root_com_ang_vel_w_ta @property - def body_link_pos_w(self) -> wp.array: + def body_link_pos_w(self) -> ProxyArray: """Positions of all bodies in simulation world frame. Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to (num_instances, num_bodies, 3). This quantity is the position of the articulation bodies' actor frame relative to the world. """ - return self._get_pos_from_transform(self.body_link_pose_w) + parent = self.body_link_pose_w + if self._body_link_pos_w_ta is None: + self._body_link_pos_w_ta = ProxyArray(self._get_pos_from_transform(parent.warp)) + return self._body_link_pos_w_ta @property - def body_link_quat_w(self) -> wp.array: + def body_link_quat_w(self) -> ProxyArray: """Orientation (x, y, z, w) of all bodies in simulation world frame. Shape is (num_instances, num_bodies), dtype = wp.quatf. In torch this resolves to (num_instances, num_bodies, 4). This quantity is the orientation of the articulation bodies' actor frame relative to the world. """ - return self._get_quat_from_transform(self.body_link_pose_w) + parent = self.body_link_pose_w + if self._body_link_quat_w_ta is None: + self._body_link_quat_w_ta = ProxyArray(self._get_quat_from_transform(parent.warp)) + return self._body_link_quat_w_ta @property - def body_link_lin_vel_w(self) -> wp.array: + def body_link_lin_vel_w(self) -> ProxyArray: """Linear velocity of all bodies in simulation world frame [m/s]. Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to (num_instances, num_bodies, 3). This quantity is the linear velocity of the articulation bodies' actor frame relative to the world. """ - return self._get_lin_vel_from_spatial_vector(self.body_link_vel_w) + parent = self.body_link_vel_w + if self._body_link_lin_vel_w_ta is None: + self._body_link_lin_vel_w_ta = ProxyArray(self._get_lin_vel_from_spatial_vector(parent.warp)) + return self._body_link_lin_vel_w_ta @property - def body_link_ang_vel_w(self) -> wp.array: + def body_link_ang_vel_w(self) -> ProxyArray: """Angular velocity of all bodies in simulation world frame [rad/s]. Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to (num_instances, num_bodies, 3). This quantity is the angular velocity of the articulation bodies' actor frame relative to the world. """ - return self._get_ang_vel_from_spatial_vector(self.body_link_vel_w) + parent = self.body_link_vel_w + if self._body_link_ang_vel_w_ta is None: + self._body_link_ang_vel_w_ta = ProxyArray(self._get_ang_vel_from_spatial_vector(parent.warp)) + return self._body_link_ang_vel_w_ta @property - def body_com_pos_w(self) -> wp.array: + def body_com_pos_w(self) -> ProxyArray: """Positions of all bodies' center of mass in simulation world frame. Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to (num_instances, num_bodies, 3). This quantity is the position of the articulation bodies' center of mass frame. """ - return self._get_pos_from_transform(self.body_com_pose_w) + parent = self.body_com_pose_w + if self._body_com_pos_w_ta is None: + self._body_com_pos_w_ta = ProxyArray(self._get_pos_from_transform(parent.warp)) + return self._body_com_pos_w_ta @property - def body_com_quat_w(self) -> wp.array: + def body_com_quat_w(self) -> ProxyArray: """Orientation (x, y, z, w) of the principal axes of inertia of all bodies in simulation world frame. Shape is (num_instances, num_bodies), dtype = wp.quatf. In torch this resolves to (num_instances, num_bodies, 4). This quantity is the orientation of the articulation bodies' principal axes of inertia. """ - return self._get_quat_from_transform(self.body_com_pose_w) + parent = self.body_com_pose_w + if self._body_com_quat_w_ta is None: + self._body_com_quat_w_ta = ProxyArray(self._get_quat_from_transform(parent.warp)) + return self._body_com_quat_w_ta @property - def body_com_lin_vel_w(self) -> wp.array: + def body_com_lin_vel_w(self) -> ProxyArray: """Linear velocity of all bodies in simulation world frame [m/s]. Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to (num_instances, num_bodies, 3). This quantity is the linear velocity of the articulation bodies' center of mass frame. """ - return self._get_lin_vel_from_spatial_vector(self.body_com_vel_w) + parent = self.body_com_vel_w + if self._body_com_lin_vel_w_ta is None: + self._body_com_lin_vel_w_ta = ProxyArray(self._get_lin_vel_from_spatial_vector(parent.warp)) + return self._body_com_lin_vel_w_ta @property - def body_com_ang_vel_w(self) -> wp.array: + def body_com_ang_vel_w(self) -> ProxyArray: """Angular velocity of all bodies in simulation world frame [rad/s]. Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to (num_instances, num_bodies, 3). This quantity is the angular velocity of the articulation bodies' center of mass frame. """ - return self._get_ang_vel_from_spatial_vector(self.body_com_vel_w) + parent = self.body_com_vel_w + if self._body_com_ang_vel_w_ta is None: + self._body_com_ang_vel_w_ta = ProxyArray(self._get_ang_vel_from_spatial_vector(parent.warp)) + return self._body_com_ang_vel_w_ta @property - def body_com_lin_acc_w(self) -> wp.array: + def body_com_lin_acc_w(self) -> ProxyArray: """Linear acceleration of all bodies in simulation world frame [m/s^2]. Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to (num_instances, num_bodies, 3). This quantity is the linear acceleration of the articulation bodies' center of mass frame. """ - return self._get_lin_vel_from_spatial_vector(self.body_com_acc_w) + parent = self.body_com_acc_w + if self._body_com_lin_acc_w_ta is None: + self._body_com_lin_acc_w_ta = ProxyArray(self._get_lin_vel_from_spatial_vector(parent.warp)) + return self._body_com_lin_acc_w_ta @property - def body_com_ang_acc_w(self) -> wp.array: + def body_com_ang_acc_w(self) -> ProxyArray: """Angular acceleration of all bodies in simulation world frame [rad/s^2]. Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to (num_instances, num_bodies, 3). This quantity is the angular acceleration of the articulation bodies' center of mass frame. """ - return self._get_ang_vel_from_spatial_vector(self.body_com_acc_w) + parent = self.body_com_acc_w + if self._body_com_ang_acc_w_ta is None: + self._body_com_ang_acc_w_ta = ProxyArray(self._get_ang_vel_from_spatial_vector(parent.warp)) + return self._body_com_ang_acc_w_ta @property - def body_com_pos_b(self) -> wp.array: + def body_com_pos_b(self) -> ProxyArray: """Center of mass position of all of the bodies in their respective link frames. Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to (num_instances, num_bodies, 3). This quantity is the center of mass location relative to its body's link frame. """ - return self._get_pos_from_transform(self.body_com_pose_b) + parent = self.body_com_pose_b + if self._body_com_pos_b_ta is None: + self._body_com_pos_b_ta = ProxyArray(self._get_pos_from_transform(parent.warp)) + return self._body_com_pos_b_ta @property - def body_com_quat_b(self) -> wp.array: + def body_com_quat_b(self) -> ProxyArray: """Orientation (x, y, z, w) of the principal axes of inertia of all of the bodies in their respective link frames. Shape is (num_instances, num_bodies), dtype = wp.quatf. In torch this resolves to @@ -1089,14 +1247,17 @@ def body_com_quat_b(self) -> wp.array: This quantity is the orientation of the principal axes of inertia relative to its body's link frame. """ - return self._get_quat_from_transform(self.body_com_pose_b) + parent = self.body_com_pose_b + if self._body_com_quat_b_ta is None: + self._body_com_quat_b_ta = ProxyArray(self._get_quat_from_transform(parent.warp)) + return self._body_com_quat_b_ta """ Deprecated in base class (required by ABC for backward compatibility). """ @property - def default_root_state(self) -> wp.array: + def default_root_state(self) -> ProxyArray: """Deprecated. Use :attr:`default_root_pose` and :attr:`default_root_vel` instead.""" warnings.warn( "default_root_state is deprecated. Use default_root_pose and default_root_vel.", @@ -1107,10 +1268,12 @@ def default_root_state(self) -> wp.array: self._root_state_w_buf = wp.zeros( self._num_instances, dtype=wp.types.vector(13, wp.float32), device=self.device ) - return self._root_state_w_buf + if self._default_root_state_ta is None: + self._default_root_state_ta = ProxyArray(self._root_state_w_buf) + return self._default_root_state_ta @property - def root_state_w(self) -> wp.array: + def root_state_w(self) -> ProxyArray: """Deprecated. Use :attr:`root_link_pose_w` and :attr:`root_com_vel_w` instead.""" warnings.warn( "root_state_w is deprecated. Use root_link_pose_w and root_com_vel_w.", @@ -1120,7 +1283,7 @@ def root_state_w(self) -> wp.array: return self.root_link_pose_w @property - def root_link_state_w(self) -> wp.array: + def root_link_state_w(self) -> ProxyArray: """Deprecated. Use :attr:`root_link_pose_w` and :attr:`root_link_vel_w` instead.""" warnings.warn( "root_link_state_w is deprecated. Use root_link_pose_w and root_link_vel_w.", @@ -1130,7 +1293,7 @@ def root_link_state_w(self) -> wp.array: return self.root_link_pose_w @property - def root_com_state_w(self) -> wp.array: + def root_com_state_w(self) -> ProxyArray: """Deprecated. Use :attr:`root_com_pose_w` and :attr:`root_com_vel_w` instead.""" warnings.warn( "root_com_state_w is deprecated. Use root_com_pose_w and root_com_vel_w.", @@ -1140,7 +1303,7 @@ def root_com_state_w(self) -> wp.array: return self.root_com_pose_w @property - def body_state_w(self) -> wp.array: + def body_state_w(self) -> ProxyArray: """Deprecated. Use :attr:`body_link_pose_w` and :attr:`body_com_vel_w` instead.""" warnings.warn( "body_state_w is deprecated. Use body_link_pose_w and body_com_vel_w.", @@ -1150,7 +1313,7 @@ def body_state_w(self) -> wp.array: return self.body_link_pose_w @property - def body_link_state_w(self) -> wp.array: + def body_link_state_w(self) -> ProxyArray: """Deprecated. Use :attr:`body_link_pose_w` and :attr:`body_link_vel_w` instead.""" warnings.warn( "body_link_state_w is deprecated. Use body_link_pose_w and body_link_vel_w.", @@ -1160,7 +1323,7 @@ def body_link_state_w(self) -> wp.array: return self.body_link_pose_w @property - def body_com_state_w(self) -> wp.array: + def body_com_state_w(self) -> ProxyArray: """Deprecated. Use :attr:`body_com_pose_w` and :attr:`body_com_vel_w` instead.""" warnings.warn( "body_com_state_w is deprecated. Use body_com_pose_w and body_com_vel_w.", @@ -1284,6 +1447,9 @@ def _create_buffers(self) -> None: # noqa: C901 # Read initial joint properties from bindings self._read_initial_properties() + # Initialize ProxyArray wrappers (lazily created on first access) + self._pin_proxy_arrays() + def _read_initial_properties(self) -> None: """Read static/initial joint and body properties from ovphysx bindings. @@ -1364,6 +1530,103 @@ def _read_cpu(tensor_type): if np_buf is not None and dst is not None: wp.copy(dst, wp.from_numpy(np_buf, dtype=wp.float32, device=self.device)) + def _pin_proxy_arrays(self) -> None: + """Create pinned ProxyArray wrappers for all data buffers. + + This is called once from :meth:`_create_buffers` during initialization. + All ``_ta`` fields are lazily populated on first property access. + """ + # -- Pinned ProxyArray cache (one per read property, lazily created on first access) + # Defaults + self._default_root_pose_ta: ProxyArray | None = None + self._default_root_vel_ta: ProxyArray | None = None + self._default_joint_pos_ta: ProxyArray | None = None + self._default_joint_vel_ta: ProxyArray | None = None + # Joint commands (set into simulation) + self._joint_pos_target_ta: ProxyArray | None = None + self._joint_vel_target_ta: ProxyArray | None = None + self._joint_effort_target_ta: ProxyArray | None = None + # Joint commands (explicit actuator model) + self._computed_torque_ta: ProxyArray | None = None + self._applied_torque_ta: ProxyArray | None = None + # Joint properties + self._joint_stiffness_ta: ProxyArray | None = None + self._joint_damping_ta: ProxyArray | None = None + self._joint_armature_ta: ProxyArray | None = None + self._joint_friction_coeff_ta: ProxyArray | None = None + self._joint_pos_limits_ta: ProxyArray | None = None + self._joint_vel_limits_ta: ProxyArray | None = None + self._joint_effort_limits_ta: ProxyArray | None = None + # Joint properties (custom) + self._soft_joint_pos_limits_ta: ProxyArray | None = None + self._soft_joint_vel_limits_ta: ProxyArray | None = None + self._gear_ratio_ta: ProxyArray | None = None + # Fixed tendon properties + self._fixed_tendon_stiffness_ta: ProxyArray | None = None + self._fixed_tendon_damping_ta: ProxyArray | None = None + self._fixed_tendon_limit_stiffness_ta: ProxyArray | None = None + self._fixed_tendon_rest_length_ta: ProxyArray | None = None + self._fixed_tendon_offset_ta: ProxyArray | None = None + self._fixed_tendon_pos_limits_ta: ProxyArray | None = None + # Spatial tendon properties + self._spatial_tendon_stiffness_ta: ProxyArray | None = None + self._spatial_tendon_damping_ta: ProxyArray | None = None + self._spatial_tendon_limit_stiffness_ta: ProxyArray | None = None + self._spatial_tendon_offset_ta: ProxyArray | None = None + # Root state (timestamped) + self._root_link_pose_w_ta: ProxyArray | None = None + self._root_link_vel_w_ta: ProxyArray | None = None + self._root_com_pose_w_ta: ProxyArray | None = None + self._root_com_vel_w_ta: ProxyArray | None = None + # Body state (timestamped) + self._body_link_pose_w_ta: ProxyArray | None = None + self._body_link_vel_w_ta: ProxyArray | None = None + self._body_com_pose_w_ta: ProxyArray | None = None + self._body_com_acc_w_ta: ProxyArray | None = None + self._body_com_pose_b_ta: ProxyArray | None = None + self._body_incoming_joint_wrench_b_ta: ProxyArray | None = None + # Body properties + self._body_mass_ta: ProxyArray | None = None + self._body_inertia_ta: ProxyArray | None = None + # Joint state (timestamped) + self._joint_pos_ta: ProxyArray | None = None + self._joint_vel_ta: ProxyArray | None = None + self._joint_acc_ta: ProxyArray | None = None + # Derived properties (timestamped) + self._projected_gravity_b_ta: ProxyArray | None = None + self._heading_w_ta: ProxyArray | None = None + self._root_link_lin_vel_b_ta: ProxyArray | None = None + self._root_link_ang_vel_b_ta: ProxyArray | None = None + self._root_com_lin_vel_b_ta: ProxyArray | None = None + self._root_com_ang_vel_b_ta: ProxyArray | None = None + # Sliced properties (root link) + self._root_link_pos_w_ta: ProxyArray | None = None + self._root_link_quat_w_ta: ProxyArray | None = None + self._root_link_lin_vel_w_ta: ProxyArray | None = None + self._root_link_ang_vel_w_ta: ProxyArray | None = None + # Sliced properties (root com) + self._root_com_pos_w_ta: ProxyArray | None = None + self._root_com_quat_w_ta: ProxyArray | None = None + self._root_com_lin_vel_w_ta: ProxyArray | None = None + self._root_com_ang_vel_w_ta: ProxyArray | None = None + # Sliced properties (body link) + self._body_link_pos_w_ta: ProxyArray | None = None + self._body_link_quat_w_ta: ProxyArray | None = None + self._body_link_lin_vel_w_ta: ProxyArray | None = None + self._body_link_ang_vel_w_ta: ProxyArray | None = None + # Sliced properties (body com) + self._body_com_pos_w_ta: ProxyArray | None = None + self._body_com_quat_w_ta: ProxyArray | None = None + self._body_com_lin_vel_w_ta: ProxyArray | None = None + self._body_com_ang_vel_w_ta: ProxyArray | None = None + self._body_com_lin_acc_w_ta: ProxyArray | None = None + self._body_com_ang_acc_w_ta: ProxyArray | None = None + # Sliced properties (body com in body frame) + self._body_com_pos_b_ta: ProxyArray | None = None + self._body_com_quat_b_ta: ProxyArray | None = None + # Deprecated state-concat properties + self._default_root_state_ta: ProxyArray | None = None + """ Internal helpers -- Bindings. """ diff --git a/source/isaaclab_ovphysx/test/assets/test_articulation_data.py b/source/isaaclab_ovphysx/test/assets/test_articulation_data.py index de17b611eda4..16bb99a4d6c4 100644 --- a/source/isaaclab_ovphysx/test/assets/test_articulation_data.py +++ b/source/isaaclab_ovphysx/test/assets/test_articulation_data.py @@ -35,7 +35,7 @@ def test_joint_acc_uses_inverse_dt(self): data.update(dt=0.25) np.testing.assert_allclose( - data.joint_acc.numpy(), + data.joint_acc.warp.numpy(), np.array([[4.0, -8.0]], dtype=np.float32), atol=1e-6, err_msg="Joint acceleration should be computed as delta_velocity / dt.", diff --git a/source/isaaclab_physx/config/extension.toml b/source/isaaclab_physx/config/extension.toml index 39a9b2dbe4a9..6abbecae4e0c 100644 --- a/source/isaaclab_physx/config/extension.toml +++ b/source/isaaclab_physx/config/extension.toml @@ -1,7 +1,7 @@ [package] # Note: Semantic Versioning is used: https://semver.org/ -version = "0.5.25" +version = "0.5.27" # 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 be7e0184817b..35d18225cf50 100644 --- a/source/isaaclab_physx/docs/CHANGELOG.rst +++ b/source/isaaclab_physx/docs/CHANGELOG.rst @@ -1,7 +1,27 @@ Changelog --------- -0.5.25 (2026-04-24) +0.5.27 (2026-04-27) +~~~~~~~~~~~~~~~~~~~ + +Fixed +^^^^^ + +* Fixed ``Simulation view object is invalidated and cannot be used again to call + getDofVelocities`` raised on the first ``scene.update()`` after ``sim.reset()`` + with recent Isaac Sim ``develop`` builds. Isaac Sim's + ``isaacsim.core.simulation_manager.SimulationManager`` recently became reactive + to timeline ``STOP`` events (after its ``_on_stop`` was decorated with + ``@staticmethod`` upstream), and its ``invalidate_physics()`` was clobbering + the shared ``omni.physics.tensors`` simulation view that + :class:`~isaaclab_physx.physics.PhysxManager` and PhysX articulation views + rely on. The ``isaaclab_physx`` package init now disables the original Isaac + Sim ``SimulationManager``'s default timeline/stage callbacks via + ``enable_all_default_callbacks(False)`` before swapping the module attribute, + so :class:`PhysxManager` is the single owner of the simulation lifecycle. + + +0.5.26 (2026-04-27) ~~~~~~~~~~~~~~~~~~~ Fixed @@ -24,7 +44,7 @@ Changed pattern used by the Kit perspective video helper. -0.5.24 (2026-04-22) +0.5.25 (2026-04-27) ~~~~~~~~~~~~~~~~~~~ Changed @@ -41,7 +61,39 @@ Changed Isaac Sim module paths (``isaacsim.core.experimental.utils.app`` and ``isaacsim.core.rendering_manager``). -0.5.23 (2026-04-23) +0.5.24 (2026-04-24) +~~~~~~~~~~~~~~~~~~~ + +Changed +^^^^^^^ + +* Updated :class:`~isaaclab_physx.sim.views.FabricFrameView` to match the new + :class:`~isaaclab.sim.views.BaseFrameView` ProxyArray return contract. See + the ``isaaclab`` 4.6.15 changelog for migration guidance. + + +0.5.23 (2026-04-24) +~~~~~~~~~~~~~~~~~~~ + +Changed +^^^^^^^ + +* Properties on the following data classes now return + :class:`~isaaclab.utils.warp.ProxyArray` instead of raw ``wp.array``: + :class:`~isaaclab_physx.assets.articulation.ArticulationData`, + :class:`~isaaclab_physx.assets.rigid_object.RigidObjectData`, + :class:`~isaaclab_physx.assets.rigid_object_collection.RigidObjectCollectionData`, + :class:`~isaaclab_physx.assets.deformable_object.DeformableObjectData`, + :class:`~isaaclab_physx.sensors.contact_sensor.ContactSensorData`, + :class:`~isaaclab_physx.sensors.frame_transformer.FrameTransformerData`, + :class:`~isaaclab_physx.sensors.imu.ImuData`, and + :class:`~isaaclab_physx.sensors.pva.PvaData`. + Use ``.torch`` for a cached zero-copy ``torch.Tensor`` view, or ``.warp`` for + the underlying ``wp.array``. Implicit torch operations (arithmetic, + ``torch.*`` functions) work during the deprecation period but emit a warning. + + +0.5.22 (2026-04-23) ~~~~~~~~~~~~~~~~~~~ Fixed @@ -55,7 +107,7 @@ Fixed the freshly built artifact on the simulation context so subsequent providers reuse it. -0.5.22 (2026-04-22) +0.5.21 (2026-04-22) ~~~~~~~~~~~~~~~~~~~ Added @@ -71,26 +123,6 @@ Changed :class:`~isaaclab_physx.sim.views.FabricFrameView`. Old name is kept as a deprecated alias. -0.5.21 (2026-04-22) -~~~~~~~~~~~~~~~~~~~ - -Fixed -^^^^^ - -* Fixed ``Simulation view object is invalidated and cannot be used again to call - getDofVelocities`` raised on the first ``scene.update()`` after ``sim.reset()`` - with recent Isaac Sim ``develop`` builds. Isaac Sim's - ``isaacsim.core.simulation_manager.SimulationManager`` recently became reactive - to timeline ``STOP`` events (after its ``_on_stop`` was decorated with - ``@staticmethod`` upstream), and its ``invalidate_physics()`` was clobbering - the shared ``omni.physics.tensors`` simulation view that - :class:`~isaaclab_physx.physics.PhysxManager` and PhysX articulation views - rely on. The ``isaaclab_physx`` package init now disables the original Isaac - Sim ``SimulationManager``'s default timeline/stage callbacks via - ``enable_all_default_callbacks(False)`` before swapping the module attribute, - so :class:`PhysxManager` is the single owner of the simulation lifecycle. - - 0.5.20 (2026-04-21) ~~~~~~~~~~~~~~~~~~~ diff --git a/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py b/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py index 8ed288a4d9d2..f2459728de98 100644 --- a/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py +++ b/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py @@ -243,8 +243,8 @@ def write_data_to_sim(self): composer = self._permanent_wrench_composer composer.compose_to_body_frame() self.root_view.apply_forces_and_torques_at_position( - force_data=composer.out_force_b.flatten().view(wp.float32), - torque_data=composer.out_torque_b.flatten().view(wp.float32), + force_data=composer.out_force_b.warp.flatten().view(wp.float32), + torque_data=composer.out_torque_b.warp.flatten().view(wp.float32), position_data=None, indices=self._ALL_INDICES, is_global=False, @@ -3076,12 +3076,12 @@ def write_fixed_tendon_properties_to_sim_index( env_ids = self._resolve_env_ids(env_ids) # Write fixed tendon properties to the simulation. self.root_view.set_fixed_tendon_properties( - self.data.fixed_tendon_stiffness, - self.data.fixed_tendon_damping, - self.data.fixed_tendon_limit_stiffness, - self.data.fixed_tendon_pos_limits, - self.data.fixed_tendon_rest_length, - self.data.fixed_tendon_offset, + self.data.fixed_tendon_stiffness.warp, + self.data.fixed_tendon_damping.warp, + self.data.fixed_tendon_limit_stiffness.warp, + self.data.fixed_tendon_pos_limits.warp, + self.data.fixed_tendon_rest_length.warp, + self.data.fixed_tendon_offset.warp, indices=env_ids, ) @@ -3533,10 +3533,10 @@ def write_spatial_tendon_properties_to_sim_index( env_ids = wp.array(env_ids, dtype=wp.int32, device=self.device) # Write spatial tendon properties to the simulation. self.root_view.set_spatial_tendon_properties( - self.data.spatial_tendon_stiffness, - self.data.spatial_tendon_damping, - self.data.spatial_tendon_limit_stiffness, - self.data.spatial_tendon_offset, + self.data.spatial_tendon_stiffness.warp, + self.data.spatial_tendon_damping.warp, + self.data.spatial_tendon_limit_stiffness.warp, + self.data.spatial_tendon_offset.warp, indices=env_ids, ) @@ -3756,14 +3756,14 @@ def _process_actuators_cfg(self): joint_ids=joint_ids, num_envs=self.num_instances, device=self.device, - stiffness=wp.to_torch(self._data.joint_stiffness)[:, joint_ids], - damping=wp.to_torch(self._data.joint_damping)[:, joint_ids], - armature=wp.to_torch(self._data.joint_armature)[:, joint_ids], - friction=wp.to_torch(self._data.joint_friction_coeff)[:, joint_ids], - dynamic_friction=wp.to_torch(self._data.joint_dynamic_friction_coeff)[:, joint_ids], - viscous_friction=wp.to_torch(self._data.joint_viscous_friction_coeff)[:, joint_ids], - effort_limit=wp.to_torch(self._data.joint_effort_limits)[:, joint_ids].clone(), - velocity_limit=wp.to_torch(self._data.joint_vel_limits)[:, joint_ids], + stiffness=self._data.joint_stiffness.torch[:, joint_ids], + damping=self._data.joint_damping.torch[:, joint_ids], + armature=self._data.joint_armature.torch[:, joint_ids], + friction=self._data.joint_friction_coeff.torch[:, joint_ids], + dynamic_friction=self._data.joint_dynamic_friction_coeff.torch[:, joint_ids], + viscous_friction=self._data.joint_viscous_friction_coeff.torch[:, joint_ids], + effort_limit=self._data.joint_effort_limits.torch[:, joint_ids].clone(), + velocity_limit=self._data.joint_vel_limits.torch[:, joint_ids], ) # store actuator group self.actuators[actuator_name] = actuator @@ -3945,16 +3945,16 @@ def _apply_actuator_model(self): # prepare input for actuator model based on cached data # TODO : A tensor dict would be nice to do the indexing of all tensors together control_action = ArticulationActions( - joint_positions=wp.to_torch(self._data.joint_pos_target)[:, actuator.joint_indices], - joint_velocities=wp.to_torch(self._data.joint_vel_target)[:, actuator.joint_indices], - joint_efforts=wp.to_torch(self._data.joint_effort_target)[:, actuator.joint_indices], + joint_positions=self._data.joint_pos_target.torch[:, actuator.joint_indices], + joint_velocities=self._data.joint_vel_target.torch[:, actuator.joint_indices], + joint_efforts=self._data.joint_effort_target.torch[:, actuator.joint_indices], joint_indices=actuator.joint_indices, ) # compute joint command from the actuator model control_action = actuator.compute( control_action, - joint_pos=wp.to_torch(self._data.joint_pos)[:, actuator.joint_indices], - joint_vel=wp.to_torch(self._data.joint_vel)[:, actuator.joint_indices], + joint_pos=self._data.joint_pos.torch[:, actuator.joint_indices], + joint_vel=self._data.joint_vel.torch[:, actuator.joint_indices], ) # update targets (these are set into the simulation) joint_indices = actuator.joint_indices @@ -4014,8 +4014,8 @@ def _validate_cfg(self): """ # check that the default values are within the limits joint_pos_limits = wp.to_torch(wp.clone(self.root_view.get_dof_limits(), device=self.device))[0] - out_of_range = wp.to_torch(self._data.default_joint_pos)[0] < joint_pos_limits[:, 0] - out_of_range |= wp.to_torch(self._data.default_joint_pos)[0] > joint_pos_limits[:, 1] + out_of_range = self._data.default_joint_pos.torch[0] < joint_pos_limits[:, 0] + out_of_range |= self._data.default_joint_pos.torch[0] > joint_pos_limits[:, 1] violated_indices = torch.nonzero(out_of_range, as_tuple=False).squeeze(-1) # throw error if any of the default joint positions are out of the limits if len(violated_indices) > 0: @@ -4024,14 +4024,14 @@ def _validate_cfg(self): for idx in violated_indices: joint_name = self.data.joint_names[idx] joint_limit = joint_pos_limits[idx] - joint_pos = wp.to_torch(self._data.default_joint_pos)[0, idx] + joint_pos = self._data.default_joint_pos.torch[0, idx] # add to message msg += f"\t- '{joint_name}': {joint_pos:.3f} not in [{joint_limit[0]:.3f}, {joint_limit[1]:.3f}]\n" raise ValueError(msg) # check that the default joint velocities are within the limits joint_max_vel = wp.to_torch(wp.clone(self.root_view.get_dof_max_velocities(), device=self.device))[0] - out_of_range = torch.abs(wp.to_torch(self._data.default_joint_vel)[0]) > joint_max_vel + out_of_range = torch.abs(self._data.default_joint_vel.torch[0]) > joint_max_vel violated_indices = torch.nonzero(out_of_range, as_tuple=False).squeeze(-1) if len(violated_indices) > 0: # prepare message for violated joints @@ -4039,7 +4039,7 @@ def _validate_cfg(self): for idx in violated_indices: joint_name = self.data.joint_names[idx] joint_limit = [-joint_max_vel[idx], joint_max_vel[idx]] - joint_vel = wp.to_torch(self._data.default_joint_vel)[0, idx] + joint_vel = self._data.default_joint_vel.torch[0, idx] # add to message msg += f"\t- '{joint_name}': {joint_vel:.3f} not in [{joint_limit[0]:.3f}, {joint_limit[1]:.3f}]\n" raise ValueError(msg) @@ -4070,23 +4070,23 @@ def format_limits(_, v: tuple[float, float]) -> str: # -- gains # Use data properties which have already been cloned and stored during initialization # This avoids issues with indexedarray or empty arrays from root_view - stiffnesses = wp.to_torch(self.data.joint_stiffness)[0].cpu().tolist() - dampings = wp.to_torch(self.data.joint_damping)[0].cpu().tolist() + stiffnesses = self.data.joint_stiffness.torch[0].cpu().tolist() + dampings = self.data.joint_damping.torch[0].cpu().tolist() # -- properties - armatures = wp.to_torch(self.data.joint_armature)[0].cpu().tolist() + armatures = self.data.joint_armature.torch[0].cpu().tolist() # For friction, use the individual components from data - friction_coeff = wp.to_torch(self.data.joint_friction_coeff)[0].cpu() - dynamic_friction_coeff = wp.to_torch(self.data.joint_dynamic_friction_coeff)[0].cpu() - viscous_friction_coeff = wp.to_torch(self.data.joint_viscous_friction_coeff)[0].cpu() + friction_coeff = self.data.joint_friction_coeff.torch[0].cpu() + dynamic_friction_coeff = self.data.joint_dynamic_friction_coeff.torch[0].cpu() + viscous_friction_coeff = self.data.joint_viscous_friction_coeff.torch[0].cpu() static_frictions = friction_coeff.tolist() dynamic_frictions = dynamic_friction_coeff.tolist() viscous_frictions = viscous_friction_coeff.tolist() # -- limits # joint_pos_limits is vec2f array, convert to torch and extract [lower, upper] pairs - position_limits_torch = wp.to_torch(self.data.joint_pos_limits)[0].cpu() # shape: (num_joints, 2) + position_limits_torch = self.data.joint_pos_limits.torch[0].cpu() # shape: (num_joints, 2) position_limits = [tuple(pos_limit.tolist()) for pos_limit in position_limits_torch] - velocity_limits = wp.to_torch(self.data.joint_vel_limits)[0].cpu().tolist() - effort_limits = wp.to_torch(self.data.joint_effort_limits)[0].cpu().tolist() + velocity_limits = self.data.joint_vel_limits.torch[0].cpu().tolist() + effort_limits = self.data.joint_effort_limits.torch[0].cpu().tolist() # create table for term information joint_table = PrettyTable() joint_table.title = f"Simulation Joint Information (Prim path: {self.cfg.prim_path})" @@ -4127,15 +4127,15 @@ def format_limits(_, v: tuple[float, float]) -> str: if self.num_fixed_tendons > 0: # -- gains # Use data properties which have already been cloned and stored during initialization - ft_stiffnesses = wp.to_torch(self.data.fixed_tendon_stiffness)[0].cpu().tolist() - ft_dampings = wp.to_torch(self.data.fixed_tendon_damping)[0].cpu().tolist() + ft_stiffnesses = self.data.fixed_tendon_stiffness.torch[0].cpu().tolist() + ft_dampings = self.data.fixed_tendon_damping.torch[0].cpu().tolist() # -- limits - ft_limit_stiffnesses = wp.to_torch(self.data.fixed_tendon_limit_stiffness)[0].cpu().tolist() + ft_limit_stiffnesses = self.data.fixed_tendon_limit_stiffness.torch[0].cpu().tolist() # fixed_tendon_pos_limits is vec2f array - ft_limits_torch = wp.to_torch(self.data.fixed_tendon_pos_limits)[0].cpu() + ft_limits_torch = self.data.fixed_tendon_pos_limits.torch[0].cpu() ft_limits = [tuple(limit.tolist()) for limit in ft_limits_torch] - ft_rest_lengths = wp.to_torch(self.data.fixed_tendon_rest_length)[0].cpu().tolist() - ft_offsets = wp.to_torch(self.data.fixed_tendon_offset)[0].cpu().tolist() + ft_rest_lengths = self.data.fixed_tendon_rest_length.torch[0].cpu().tolist() + ft_offsets = self.data.fixed_tendon_offset.torch[0].cpu().tolist() # create table for term information tendon_table = PrettyTable() tendon_table.title = f"Simulation Fixed Tendon Information (Prim path: {self.cfg.prim_path})" @@ -4179,11 +4179,11 @@ def format_limits(_, v: tuple[float, float]) -> str: if self.num_spatial_tendons > 0: # -- gains # Use data properties which have already been cloned and stored during initialization - st_stiffnesses = wp.to_torch(self.data.spatial_tendon_stiffness)[0].cpu().tolist() - st_dampings = wp.to_torch(self.data.spatial_tendon_damping)[0].cpu().tolist() + st_stiffnesses = self.data.spatial_tendon_stiffness.torch[0].cpu().tolist() + st_dampings = self.data.spatial_tendon_damping.torch[0].cpu().tolist() # -- limits - st_limit_stiffnesses = wp.to_torch(self.data.spatial_tendon_limit_stiffness)[0].cpu().tolist() - st_offsets = wp.to_torch(self.data.spatial_tendon_offset)[0].cpu().tolist() + st_limit_stiffnesses = self.data.spatial_tendon_limit_stiffness.torch[0].cpu().tolist() + st_offsets = self.data.spatial_tendon_offset.torch[0].cpu().tolist() # create table for term information tendon_table = PrettyTable() tendon_table.title = f"Simulation Spatial Tendon Information (Prim path: {self.cfg.prim_path})" 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 b9955c042785..0e27e05bd141 100644 --- a/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation_data.py +++ b/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation_data.py @@ -15,6 +15,7 @@ from isaaclab.assets.articulation.base_articulation_data import BaseArticulationData from isaaclab.utils.buffers import TimestampedBufferWarp as TimestampedBuffer from isaaclab.utils.math import normalize +from isaaclab.utils.warp import ProxyArray from isaaclab_physx.assets import kernels as shared_kernels from isaaclab_physx.assets.articulation import kernels as articulation_kernels @@ -43,6 +44,21 @@ class ArticulationData(BaseArticulationData): Depending on the settings, the two frames may not coincide with each other. In the robotics sense, the actor frame can be interpreted as the link frame. + + .. note:: + **Pull-to-refresh model.** PhysX state properties are *not* automatically updated each + simulation step. Each property getter pulls fresh data from the PhysX tensor API on first + access per timestamp, then caches the result until the next step. This differs from the + Newton backend, where buffers are refreshed automatically by the simulation. + + .. note:: + **ProxyArray pointer stability.** Each :class:`ProxyArray` wrapper is created once on the + first property access and reused thereafter. This is safe because the PhysX tensor API + returns views into stable, pre-allocated GPU buffers whose device pointer does not change + across simulation steps. The ``wp.array`` Python objects returned by getters like + ``get_root_transforms()`` are new wrappers each call, but they alias the same underlying + GPU memory. Sub-view properties (``root_pos_w``, ``root_quat_w``, etc.) similarly wrap + pointer offsets into these stable buffers and are therefore also safe to cache. """ __backend_name__: str = "physx" @@ -75,8 +91,8 @@ def __init__(self, root_view: physx.ArticulationView, device: str): forward_vec = torch.tensor((1.0, 0.0, 0.0), device=self.device).repeat(self._root_view.count, 1) # Initialize constants - self.GRAVITY_VEC_W = wp.from_torch(gravity_dir, dtype=wp.vec3f) - self.FORWARD_VEC_B = wp.from_torch(forward_vec, dtype=wp.vec3f) + self.GRAVITY_VEC_W = ProxyArray(wp.from_torch(gravity_dir, dtype=wp.vec3f)) + self.FORWARD_VEC_B = ProxyArray(wp.from_torch(forward_vec, dtype=wp.vec3f)) self._create_buffers() @@ -135,13 +151,15 @@ def update(self, dt: float) -> None: """ @property - def default_root_pose(self) -> wp.array: + def default_root_pose(self) -> ProxyArray: """Default root pose ``[pos, quat]`` in the local environment frame. The position and quaternion are of the articulation root's actor frame. Shape is (num_instances,), dtype = wp.transformf. In torch this resolves to (num_instances, 7). """ - return self._default_root_pose + if self._default_root_pose_ta is None: + self._default_root_pose_ta = ProxyArray(self._default_root_pose) + return self._default_root_pose_ta @default_root_pose.setter def default_root_pose(self, value: wp.array) -> None: @@ -158,13 +176,15 @@ def default_root_pose(self, value: wp.array) -> None: self._default_root_pose.assign(value) @property - def default_root_vel(self) -> wp.array: + def default_root_vel(self) -> ProxyArray: """Default root velocity ``[lin_vel, ang_vel]`` in the local environment frame. The linear and angular velocities are of the articulation root's center of mass frame. Shape is (num_instances,), dtype = wp.spatial_vectorf. In torch this resolves to (num_instances, 6). """ - return self._default_root_vel + if self._default_root_vel_ta is None: + self._default_root_vel_ta = ProxyArray(self._default_root_vel) + return self._default_root_vel_ta @default_root_vel.setter def default_root_vel(self, value: wp.array) -> None: @@ -181,14 +201,16 @@ def default_root_vel(self, value: wp.array) -> None: self._default_root_vel.assign(value) @property - def default_joint_pos(self) -> wp.array: + def default_joint_pos(self) -> ProxyArray: """Default joint positions of all joints. Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to (num_instances, num_joints). This quantity is configured through the :attr:`isaaclab.assets.ArticulationCfg.init_state` parameter. """ - return self._default_joint_pos + if self._default_joint_pos_ta is None: + self._default_joint_pos_ta = ProxyArray(self._default_joint_pos) + return self._default_joint_pos_ta @default_joint_pos.setter def default_joint_pos(self, value: wp.array) -> None: @@ -205,14 +227,16 @@ def default_joint_pos(self, value: wp.array) -> None: self._default_joint_pos.assign(value) @property - def default_joint_vel(self) -> wp.array: + def default_joint_vel(self) -> ProxyArray: """Default joint velocities of all joints. Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to (num_instances, num_joints). This quantity is configured through the :attr:`isaaclab.assets.ArticulationCfg.init_state` parameter. """ - return self._default_joint_vel + if self._default_joint_vel_ta is None: + self._default_joint_vel_ta = ProxyArray(self._default_joint_vel) + return self._default_joint_vel_ta @default_joint_vel.setter def default_joint_vel(self, value: wp.array) -> None: @@ -233,7 +257,7 @@ def default_joint_vel(self, value: wp.array) -> None: """ @property - def joint_pos_target(self) -> wp.array: + def joint_pos_target(self) -> ProxyArray: """Joint position targets commanded by the user. Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to (num_instances, num_joints). @@ -242,10 +266,12 @@ def joint_pos_target(self) -> wp.array: For an explicit actuator model, the targets are used to compute the joint torques (see :attr:`applied_torque`), which are then set into the simulation. """ - return self._joint_pos_target + if self._joint_pos_target_ta is None: + self._joint_pos_target_ta = ProxyArray(self._joint_pos_target) + return self._joint_pos_target_ta @property - def joint_vel_target(self) -> wp.array: + def joint_vel_target(self) -> ProxyArray: """Joint velocity targets commanded by the user. Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to (num_instances, num_joints). @@ -254,10 +280,12 @@ def joint_vel_target(self) -> wp.array: For an explicit actuator model, the targets are used to compute the joint torques (see :attr:`applied_torque`), which are then set into the simulation. """ - return self._joint_vel_target + if self._joint_vel_target_ta is None: + self._joint_vel_target_ta = ProxyArray(self._joint_vel_target) + return self._joint_vel_target_ta @property - def joint_effort_target(self) -> wp.array: + def joint_effort_target(self) -> ProxyArray: """Joint effort targets commanded by the user. Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to (num_instances, num_joints). @@ -266,14 +294,16 @@ def joint_effort_target(self) -> wp.array: For an explicit actuator model, the targets are used to compute the joint torques (see :attr:`applied_torque`), which are then set into the simulation. """ - return self._joint_effort_target + if self._joint_effort_target_ta is None: + self._joint_effort_target_ta = ProxyArray(self._joint_effort_target) + return self._joint_effort_target_ta """ Joint commands -- Explicit actuators. """ @property - def computed_torque(self) -> wp.array: + def computed_torque(self) -> ProxyArray: """Joint torques computed from the actuator model (before clipping). Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to (num_instances, num_joints). @@ -282,10 +312,12 @@ def computed_torque(self) -> wp.array: It is exposed for users who want to inspect the computations inside the actuator model. For instance, to penalize the learning agent for a difference between the computed and applied torques. """ - return self._computed_torque + if self._computed_torque_ta is None: + self._computed_torque_ta = ProxyArray(self._computed_torque) + return self._computed_torque_ta @property - def applied_torque(self) -> wp.array: + def applied_torque(self) -> ProxyArray: """Joint torques applied from the actuator model (after clipping). Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to (num_instances, num_joints). @@ -293,66 +325,80 @@ def applied_torque(self) -> wp.array: These torques are set into the simulation, after clipping the :attr:`computed_torque` based on the actuator model. """ - return self._applied_torque + if self._applied_torque_ta is None: + self._applied_torque_ta = ProxyArray(self._applied_torque) + return self._applied_torque_ta """ Joint properties """ @property - def joint_stiffness(self) -> wp.array: + def joint_stiffness(self) -> ProxyArray: """Joint stiffness provided to the simulation. Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to (num_instances, num_joints). In the case of explicit actuators, the value for the corresponding joints is zero. """ - return self._joint_stiffness + if self._joint_stiffness_ta is None: + self._joint_stiffness_ta = ProxyArray(self._joint_stiffness) + return self._joint_stiffness_ta @property - def joint_damping(self) -> wp.array: + def joint_damping(self) -> ProxyArray: """Joint damping provided to the simulation. Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to (num_instances, num_joints). In the case of explicit actuators, the value for the corresponding joints is zero. """ - return self._joint_damping + if self._joint_damping_ta is None: + self._joint_damping_ta = ProxyArray(self._joint_damping) + return self._joint_damping_ta @property - def joint_armature(self) -> wp.array: + def joint_armature(self) -> ProxyArray: """Joint armature provided to the simulation. Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to (num_instances, num_joints). """ - return self._joint_armature + if self._joint_armature_ta is None: + self._joint_armature_ta = ProxyArray(self._joint_armature) + return self._joint_armature_ta @property - def joint_friction_coeff(self) -> wp.array: + def joint_friction_coeff(self) -> ProxyArray: """Joint static friction coefficient provided to the simulation. Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to (num_instances, num_joints). """ - return self._joint_friction_coeff + if self._joint_friction_coeff_ta is None: + self._joint_friction_coeff_ta = ProxyArray(self._joint_friction_coeff) + return self._joint_friction_coeff_ta @property - def joint_dynamic_friction_coeff(self) -> wp.array: + def joint_dynamic_friction_coeff(self) -> ProxyArray: """Joint dynamic friction coefficient provided to the simulation. Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to (num_instances, num_joints). """ - return self._joint_dynamic_friction_coeff + if self._joint_dynamic_friction_coeff_ta is None: + self._joint_dynamic_friction_coeff_ta = ProxyArray(self._joint_dynamic_friction_coeff) + return self._joint_dynamic_friction_coeff_ta @property - def joint_viscous_friction_coeff(self) -> wp.array: + def joint_viscous_friction_coeff(self) -> ProxyArray: """Joint viscous friction coefficient provided to the simulation. Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to (num_instances, num_joints). """ - return self._joint_viscous_friction_coeff + if self._joint_viscous_friction_coeff_ta is None: + self._joint_viscous_friction_coeff_ta = ProxyArray(self._joint_viscous_friction_coeff) + return self._joint_viscous_friction_coeff_ta @property - def joint_pos_limits(self) -> wp.array: + def joint_pos_limits(self) -> ProxyArray: """Joint position limits provided to the simulation. Shape is (num_instances, num_joints), dtype = wp.vec2f. In torch this resolves to @@ -360,30 +406,36 @@ def joint_pos_limits(self) -> wp.array: The limits are in the order :math:`[lower, upper]`. """ - return self._joint_pos_limits + if self._joint_pos_limits_ta is None: + self._joint_pos_limits_ta = ProxyArray(self._joint_pos_limits) + return self._joint_pos_limits_ta @property - def joint_vel_limits(self) -> wp.array: + def joint_vel_limits(self) -> ProxyArray: """Joint maximum velocity provided to the simulation. Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to (num_instances, num_joints). """ - return self._joint_vel_limits + if self._joint_vel_limits_ta is None: + self._joint_vel_limits_ta = ProxyArray(self._joint_vel_limits) + return self._joint_vel_limits_ta @property - def joint_effort_limits(self) -> wp.array: + def joint_effort_limits(self) -> ProxyArray: """Joint maximum effort provided to the simulation. Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to (num_instances, num_joints). """ - return self._joint_effort_limits + if self._joint_effort_limits_ta is None: + self._joint_effort_limits_ta = ProxyArray(self._joint_effort_limits) + return self._joint_effort_limits_ta """ Joint properties - Custom. """ @property - def soft_joint_pos_limits(self) -> wp.array: + def soft_joint_pos_limits(self) -> ProxyArray: r"""Soft joint positions limits for all joints. Shape is (num_instances, num_joints), dtype = wp.vec2f. In torch this resolves to @@ -404,10 +456,12 @@ def soft_joint_pos_limits(self) -> wp.array: The soft joint position limits help specify a safety region around the joint limits. It isn't used by the simulation, but is useful for learning agents to prevent the joint positions from violating the limits. """ - return self._soft_joint_pos_limits + if self._soft_joint_pos_limits_ta is None: + self._soft_joint_pos_limits_ta = ProxyArray(self._soft_joint_pos_limits) + return self._soft_joint_pos_limits_ta @property - def soft_joint_vel_limits(self) -> wp.array: + def soft_joint_vel_limits(self) -> ProxyArray: """Soft joint velocity limits for all joints. Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to (num_instances, num_joints). @@ -415,120 +469,144 @@ def soft_joint_vel_limits(self) -> wp.array: These are obtained from the actuator model. It may differ from :attr:`joint_vel_limits` if the actuator model has a variable velocity limit model. For instance, in a variable gear ratio actuator model. """ - return self._soft_joint_vel_limits + if self._soft_joint_vel_limits_ta is None: + self._soft_joint_vel_limits_ta = ProxyArray(self._soft_joint_vel_limits) + return self._soft_joint_vel_limits_ta @property - def gear_ratio(self) -> wp.array: + def gear_ratio(self) -> ProxyArray: """Gear ratio for relating motor torques to applied Joint torques. Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to (num_instances, num_joints). """ - return self._gear_ratio + if self._gear_ratio_ta is None: + self._gear_ratio_ta = ProxyArray(self._gear_ratio) + return self._gear_ratio_ta """ Fixed tendon properties. """ @property - def fixed_tendon_stiffness(self) -> wp.array: + def fixed_tendon_stiffness(self) -> ProxyArray: """Fixed tendon stiffness provided to the simulation. Shape is (num_instances, num_fixed_tendons), dtype = wp.float32. In torch this resolves to (num_instances, num_fixed_tendons). """ - return self._fixed_tendon_stiffness + if self._fixed_tendon_stiffness_ta is None: + self._fixed_tendon_stiffness_ta = ProxyArray(self._fixed_tendon_stiffness) + return self._fixed_tendon_stiffness_ta @property - def fixed_tendon_damping(self) -> wp.array: + def fixed_tendon_damping(self) -> ProxyArray: """Fixed tendon damping provided to the simulation. Shape is (num_instances, num_fixed_tendons), dtype = wp.float32. In torch this resolves to (num_instances, num_fixed_tendons). """ - return self._fixed_tendon_damping + if self._fixed_tendon_damping_ta is None: + self._fixed_tendon_damping_ta = ProxyArray(self._fixed_tendon_damping) + return self._fixed_tendon_damping_ta @property - def fixed_tendon_limit_stiffness(self) -> wp.array: + def fixed_tendon_limit_stiffness(self) -> ProxyArray: """Fixed tendon limit stiffness provided to the simulation. Shape is (num_instances, num_fixed_tendons), dtype = wp.float32. In torch this resolves to (num_instances, num_fixed_tendons). """ - return self._fixed_tendon_limit_stiffness + if self._fixed_tendon_limit_stiffness_ta is None: + self._fixed_tendon_limit_stiffness_ta = ProxyArray(self._fixed_tendon_limit_stiffness) + return self._fixed_tendon_limit_stiffness_ta @property - def fixed_tendon_rest_length(self) -> wp.array: + def fixed_tendon_rest_length(self) -> ProxyArray: """Fixed tendon rest length provided to the simulation. Shape is (num_instances, num_fixed_tendons), dtype = wp.float32. In torch this resolves to (num_instances, num_fixed_tendons). """ - return self._fixed_tendon_rest_length + if self._fixed_tendon_rest_length_ta is None: + self._fixed_tendon_rest_length_ta = ProxyArray(self._fixed_tendon_rest_length) + return self._fixed_tendon_rest_length_ta @property - def fixed_tendon_offset(self) -> wp.array: + def fixed_tendon_offset(self) -> ProxyArray: """Fixed tendon offset provided to the simulation. Shape is (num_instances, num_fixed_tendons), dtype = wp.float32. In torch this resolves to (num_instances, num_fixed_tendons). """ - return self._fixed_tendon_offset + if self._fixed_tendon_offset_ta is None: + self._fixed_tendon_offset_ta = ProxyArray(self._fixed_tendon_offset) + return self._fixed_tendon_offset_ta @property - def fixed_tendon_pos_limits(self) -> wp.array: + def fixed_tendon_pos_limits(self) -> ProxyArray: """Fixed tendon position limits provided to the simulation. Shape is (num_instances, num_fixed_tendons), dtype = wp.vec2f. In torch this resolves to (num_instances, num_fixed_tendons, 2). """ - return self._fixed_tendon_pos_limits + if self._fixed_tendon_pos_limits_ta is None: + self._fixed_tendon_pos_limits_ta = ProxyArray(self._fixed_tendon_pos_limits) + return self._fixed_tendon_pos_limits_ta """ Spatial tendon properties. """ @property - def spatial_tendon_stiffness(self) -> wp.array: + def spatial_tendon_stiffness(self) -> ProxyArray: """Spatial tendon stiffness provided to the simulation. Shape is (num_instances, num_spatial_tendons), dtype = wp.float32. In torch this resolves to (num_instances, num_spatial_tendons). """ - return self._spatial_tendon_stiffness + if self._spatial_tendon_stiffness_ta is None: + self._spatial_tendon_stiffness_ta = ProxyArray(self._spatial_tendon_stiffness) + return self._spatial_tendon_stiffness_ta @property - def spatial_tendon_damping(self) -> wp.array: + def spatial_tendon_damping(self) -> ProxyArray: """Spatial tendon damping provided to the simulation. Shape is (num_instances, num_spatial_tendons), dtype = wp.float32. In torch this resolves to (num_instances, num_spatial_tendons). """ - return self._spatial_tendon_damping + if self._spatial_tendon_damping_ta is None: + self._spatial_tendon_damping_ta = ProxyArray(self._spatial_tendon_damping) + return self._spatial_tendon_damping_ta @property - def spatial_tendon_limit_stiffness(self) -> wp.array: + def spatial_tendon_limit_stiffness(self) -> ProxyArray: """Spatial tendon limit stiffness provided to the simulation. Shape is (num_instances, num_spatial_tendons), dtype = wp.float32. In torch this resolves to (num_instances, num_spatial_tendons). """ - return self._spatial_tendon_limit_stiffness + if self._spatial_tendon_limit_stiffness_ta is None: + self._spatial_tendon_limit_stiffness_ta = ProxyArray(self._spatial_tendon_limit_stiffness) + return self._spatial_tendon_limit_stiffness_ta @property - def spatial_tendon_offset(self) -> wp.array: + def spatial_tendon_offset(self) -> ProxyArray: """Spatial tendon offset provided to the simulation. Shape is (num_instances, num_spatial_tendons), dtype = wp.float32. In torch this resolves to (num_instances, num_spatial_tendons). """ - return self._spatial_tendon_offset + if self._spatial_tendon_offset_ta is None: + self._spatial_tendon_offset_ta = ProxyArray(self._spatial_tendon_offset) + return self._spatial_tendon_offset_ta """ Root state properties. """ @property - def root_link_pose_w(self) -> wp.array: + def root_link_pose_w(self) -> ProxyArray: """Root link pose ``[pos, quat]`` in simulation world frame. Shape is (num_instances,), dtype = wp.transformf. In torch this resolves to (num_instances, 7). @@ -540,10 +618,12 @@ def root_link_pose_w(self) -> wp.array: self._root_link_pose_w.data = self._root_view.get_root_transforms().view(wp.transformf) self._root_link_pose_w.timestamp = self._sim_timestamp - return self._root_link_pose_w.data + if self._root_link_pose_w_ta is None: + self._root_link_pose_w_ta = ProxyArray(self._root_link_pose_w.data) + return self._root_link_pose_w_ta @property - def root_link_vel_w(self) -> wp.array: + def root_link_vel_w(self) -> ProxyArray: """Root link velocity ``[lin_vel, ang_vel]`` in simulation world frame. Shape is (num_instances,), dtype = wp.spatial_vectorf. In torch this resolves to (num_instances, 6). @@ -566,10 +646,12 @@ def root_link_vel_w(self) -> wp.array: ) self._root_link_vel_w.timestamp = self._sim_timestamp - return self._root_link_vel_w.data + if self._root_link_vel_w_ta is None: + self._root_link_vel_w_ta = ProxyArray(self._root_link_vel_w.data) + return self._root_link_vel_w_ta @property - def root_com_pose_w(self) -> wp.array: + def root_com_pose_w(self) -> ProxyArray: """Root center of mass pose ``[pos, quat]`` in simulation world frame. Shape is (num_instances,), dtype = wp.transformf. In torch this resolves to (num_instances, 7). @@ -592,10 +674,12 @@ def root_com_pose_w(self) -> wp.array: ) self._root_com_pose_w.timestamp = self._sim_timestamp - return self._root_com_pose_w.data + if self._root_com_pose_w_ta is None: + self._root_com_pose_w_ta = ProxyArray(self._root_com_pose_w.data) + return self._root_com_pose_w_ta @property - def root_com_vel_w(self) -> wp.array: + def root_com_vel_w(self) -> ProxyArray: """Root center of mass velocity ``[lin_vel, ang_vel]`` in simulation world frame. Shape is (num_instances,), dtype = wp.spatial_vectorf. In torch this resolves to (num_instances, 6). @@ -606,33 +690,39 @@ def root_com_vel_w(self) -> wp.array: self._root_com_vel_w.data = self._root_view.get_root_velocities().view(wp.spatial_vectorf) self._root_com_vel_w.timestamp = self._sim_timestamp - return self._root_com_vel_w.data + if self._root_com_vel_w_ta is None: + self._root_com_vel_w_ta = ProxyArray(self._root_com_vel_w.data) + return self._root_com_vel_w_ta """ Body state properties. """ @property - def body_mass(self) -> wp.array: + def body_mass(self) -> ProxyArray: """Body mass in the world frame. Shape is (num_instances, num_bodies), dtype = wp.float32. In torch this resolves to (num_instances, num_bodies). """ self._body_mass.assign(self._root_view.get_masses()) - return self._body_mass + if self._body_mass_ta is None: + self._body_mass_ta = ProxyArray(self._body_mass) + return self._body_mass_ta @property - def body_inertia(self) -> wp.array: + def body_inertia(self) -> ProxyArray: """Flattened body inertia in the world frame. Shape is (num_instances, num_bodies, 9), dtype = wp.float32. In torch this resolves to (num_instances, num_bodies, 9). """ self._body_inertia.assign(self._root_view.get_inertias()) - return self._body_inertia + if self._body_inertia_ta is None: + self._body_inertia_ta = ProxyArray(self._body_inertia) + return self._body_inertia_ta @property - def body_link_pose_w(self) -> wp.array: + def body_link_pose_w(self) -> ProxyArray: """Body link pose ``[pos, quat]`` in simulation world frame. Shape is (num_instances, num_bodies), dtype = wp.transformf. In torch this resolves to (num_instances, num_bodies, 7). @@ -647,10 +737,12 @@ def body_link_pose_w(self) -> wp.array: self._body_link_pose_w.data = self._root_view.get_link_transforms().view(wp.transformf) self._body_link_pose_w.timestamp = self._sim_timestamp - return self._body_link_pose_w.data + if self._body_link_pose_w_ta is None: + self._body_link_pose_w_ta = ProxyArray(self._body_link_pose_w.data) + return self._body_link_pose_w_ta @property - def body_link_vel_w(self) -> wp.array: + def body_link_vel_w(self) -> ProxyArray: """Body link velocity ``[lin_vel, ang_vel]`` in simulation world frame. Shape is (num_instances, num_bodies), dtype = wp.spatial_vectorf. In torch this resolves to (num_instances, num_bodies, 6). @@ -674,10 +766,12 @@ def body_link_vel_w(self) -> wp.array: ) self._body_link_vel_w.timestamp = self._sim_timestamp - return self._body_link_vel_w.data + if self._body_link_vel_w_ta is None: + self._body_link_vel_w_ta = ProxyArray(self._body_link_vel_w.data) + return self._body_link_vel_w_ta @property - def body_com_pose_w(self) -> wp.array: + def body_com_pose_w(self) -> ProxyArray: """Body center of mass pose ``[pos, quat]`` in simulation world frame. Shape is (num_instances, num_bodies), dtype = wp.transformf. In torch this resolves to (num_instances, num_bodies, 7). @@ -700,10 +794,12 @@ def body_com_pose_w(self) -> wp.array: ) self._body_com_pose_w.timestamp = self._sim_timestamp - return self._body_com_pose_w.data + if self._body_com_pose_w_ta is None: + self._body_com_pose_w_ta = ProxyArray(self._body_com_pose_w.data) + return self._body_com_pose_w_ta @property - def body_com_vel_w(self) -> wp.array: + def body_com_vel_w(self) -> ProxyArray: """Body center of mass velocity ``[lin_vel, ang_vel]`` in simulation world frame. Shape is (num_instances, num_bodies), dtype = wp.spatial_vectorf. In torch this resolves to (num_instances, num_bodies, 6). @@ -715,10 +811,12 @@ def body_com_vel_w(self) -> wp.array: self._body_com_vel_w.data = self._root_view.get_link_velocities().view(wp.spatial_vectorf) self._body_com_vel_w.timestamp = self._sim_timestamp - return self._body_com_vel_w.data + if self._body_com_vel_w_ta is None: + self._body_com_vel_w_ta = ProxyArray(self._body_com_vel_w.data) + return self._body_com_vel_w_ta @property - def body_com_acc_w(self): + def body_com_acc_w(self) -> ProxyArray: """Acceleration of all bodies center of mass ``[lin_acc, ang_acc]``. Shape is (num_instances, num_bodies), dtype = wp.spatial_vectorf. In torch this resolves to (num_instances, num_bodies, 6). @@ -730,10 +828,12 @@ def body_com_acc_w(self): self._body_com_acc_w.data = self._root_view.get_link_accelerations().view(wp.spatial_vectorf) self._body_com_acc_w.timestamp = self._sim_timestamp - return self._body_com_acc_w.data + if self._body_com_acc_w_ta is None: + self._body_com_acc_w_ta = ProxyArray(self._body_com_acc_w.data) + return self._body_com_acc_w_ta @property - def body_com_pose_b(self) -> wp.array: + def body_com_pose_b(self) -> ProxyArray: """Center of mass pose ``[pos, quat]`` of all bodies in their respective body's link frames. Shape is (num_instances, num_bodies), dtype = wp.transformf. In torch this resolves to (num_instances, num_bodies, 7). @@ -746,10 +846,12 @@ def body_com_pose_b(self) -> wp.array: self._body_com_pose_b.data.assign(self._root_view.get_coms().view(wp.transformf)) self._body_com_pose_b.timestamp = self._sim_timestamp - return self._body_com_pose_b.data + if self._body_com_pose_b_ta is None: + self._body_com_pose_b_ta = ProxyArray(self._body_com_pose_b.data) + return self._body_com_pose_b_ta @property - def body_incoming_joint_wrench_b(self) -> wp.array: + def body_incoming_joint_wrench_b(self) -> ProxyArray: """Joint reaction wrench applied to each body through its incoming joint, expressed in that body's frame. Shape is (num_instances, num_bodies, 6). All body reaction wrenches are provided including the root body to the @@ -771,14 +873,16 @@ def body_incoming_joint_wrench_b(self) -> wp.array: wp.spatial_vectorf ) self._body_incoming_joint_wrench_b.timestamp = self._sim_timestamp - return self._body_incoming_joint_wrench_b.data + if self._body_incoming_joint_wrench_b_ta is None: + self._body_incoming_joint_wrench_b_ta = ProxyArray(self._body_incoming_joint_wrench_b.data) + return self._body_incoming_joint_wrench_b_ta """ Joint state properties. """ @property - def joint_pos(self) -> wp.array: + def joint_pos(self) -> ProxyArray: """Joint positions of all joints. Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to (num_instances, num_joints). @@ -787,10 +891,12 @@ def joint_pos(self) -> wp.array: # read data from simulation and set the buffer data and timestamp self._joint_pos.data = self._root_view.get_dof_positions() self._joint_pos.timestamp = self._sim_timestamp - return self._joint_pos.data + if self._joint_pos_ta is None: + self._joint_pos_ta = ProxyArray(self._joint_pos.data) + return self._joint_pos_ta @property - def joint_vel(self) -> wp.array: + def joint_vel(self) -> ProxyArray: """Joint velocities of all joints. Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to (num_instances, num_joints). @@ -799,10 +905,12 @@ def joint_vel(self) -> wp.array: # read data from simulation and set the buffer data and timestamp self._joint_vel.data = self._root_view.get_dof_velocities() self._joint_vel.timestamp = self._sim_timestamp - return self._joint_vel.data + if self._joint_vel_ta is None: + self._joint_vel_ta = ProxyArray(self._joint_vel.data) + return self._joint_vel_ta @property - def joint_acc(self) -> wp.array: + def joint_acc(self) -> ProxyArray: """Joint acceleration of all joints. Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to (num_instances, num_joints). @@ -824,14 +932,16 @@ def joint_acc(self) -> wp.array: device=self.device, ) self._joint_acc.timestamp = self._sim_timestamp - return self._joint_acc.data + if self._joint_acc_ta is None: + self._joint_acc_ta = ProxyArray(self._joint_acc.data) + return self._joint_acc_ta """ Derived Properties. """ @property - def projected_gravity_b(self): + def projected_gravity_b(self) -> ProxyArray: """Projection of the gravity direction on base frame. Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3).""" if self._projected_gravity_b.timestamp < self._sim_timestamp: @@ -843,10 +953,12 @@ def projected_gravity_b(self): device=self.device, ) self._projected_gravity_b.timestamp = self._sim_timestamp - return self._projected_gravity_b.data + if self._projected_gravity_b_ta is None: + self._projected_gravity_b_ta = ProxyArray(self._projected_gravity_b.data) + return self._projected_gravity_b_ta @property - def heading_w(self): + def heading_w(self) -> ProxyArray: """Yaw heading of the base frame (in radians). Shape is (num_instances,), dtype = wp.float32. .. note:: @@ -862,10 +974,12 @@ def heading_w(self): device=self.device, ) self._heading_w.timestamp = self._sim_timestamp - return self._heading_w.data + if self._heading_w_ta is None: + self._heading_w_ta = ProxyArray(self._heading_w.data) + return self._heading_w_ta @property - def root_link_lin_vel_b(self) -> wp.array: + def root_link_lin_vel_b(self) -> ProxyArray: """Root link linear velocity in base frame. Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). @@ -880,10 +994,12 @@ def root_link_lin_vel_b(self) -> wp.array: device=self.device, ) self._root_link_lin_vel_b.timestamp = self._sim_timestamp - return self._root_link_lin_vel_b.data + if self._root_link_lin_vel_b_ta is None: + self._root_link_lin_vel_b_ta = ProxyArray(self._root_link_lin_vel_b.data) + return self._root_link_lin_vel_b_ta @property - def root_link_ang_vel_b(self) -> wp.array: + def root_link_ang_vel_b(self) -> ProxyArray: """Root link angular velocity in base frame. Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). @@ -898,10 +1014,12 @@ def root_link_ang_vel_b(self) -> wp.array: device=self.device, ) self._root_link_ang_vel_b.timestamp = self._sim_timestamp - return self._root_link_ang_vel_b.data + if self._root_link_ang_vel_b_ta is None: + self._root_link_ang_vel_b_ta = ProxyArray(self._root_link_ang_vel_b.data) + return self._root_link_ang_vel_b_ta @property - def root_com_lin_vel_b(self) -> wp.array: + def root_com_lin_vel_b(self) -> ProxyArray: """Root center of mass linear velocity in base frame. Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). @@ -917,10 +1035,12 @@ def root_com_lin_vel_b(self) -> wp.array: device=self.device, ) self._root_com_lin_vel_b.timestamp = self._sim_timestamp - return self._root_com_lin_vel_b.data + if self._root_com_lin_vel_b_ta is None: + self._root_com_lin_vel_b_ta = ProxyArray(self._root_com_lin_vel_b.data) + return self._root_com_lin_vel_b_ta @property - def root_com_ang_vel_b(self) -> wp.array: + def root_com_ang_vel_b(self) -> ProxyArray: """Root center of mass angular velocity in base frame. Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). @@ -936,196 +1056,274 @@ def root_com_ang_vel_b(self) -> wp.array: device=self.device, ) self._root_com_ang_vel_b.timestamp = self._sim_timestamp - return self._root_com_ang_vel_b.data + if self._root_com_ang_vel_b_ta is None: + self._root_com_ang_vel_b_ta = ProxyArray(self._root_com_ang_vel_b.data) + return self._root_com_ang_vel_b_ta """ Sliced properties. """ @property - def root_link_pos_w(self) -> wp.array: + def root_link_pos_w(self) -> ProxyArray: """Root link position in simulation world frame. Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). This quantity is the position of the actor frame of the root rigid body relative to the world. """ - return self._get_pos_from_transform(self.root_link_pose_w) + # Access parent property to trigger its getter call (PhysX is pull-on-demand) + parent = self.root_link_pose_w + if self._root_link_pos_w_ta is None: + self._root_link_pos_w_ta = ProxyArray(self._get_pos_from_transform(parent.warp)) + return self._root_link_pos_w_ta @property - def root_link_quat_w(self) -> wp.array: + def root_link_quat_w(self) -> ProxyArray: """Root link orientation (x, y, z, w) in simulation world frame. Shape is (num_instances,), dtype = wp.quatf. In torch this resolves to (num_instances, 4). This quantity is the orientation of the actor frame of the root rigid body. """ - return self._get_quat_from_transform(self.root_link_pose_w) + # Access parent property to trigger its getter call (PhysX is pull-on-demand) + parent = self.root_link_pose_w + if self._root_link_quat_w_ta is None: + self._root_link_quat_w_ta = ProxyArray(self._get_quat_from_transform(parent.warp)) + return self._root_link_quat_w_ta @property - def root_link_lin_vel_w(self) -> wp.array: + def root_link_lin_vel_w(self) -> ProxyArray: """Root linear velocity in simulation world frame. Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). This quantity is the linear velocity of the root rigid body's actor frame relative to the world. """ - return self._get_lin_vel_from_spatial_vector(self.root_link_vel_w) + # Access parent property to trigger its getter call (PhysX is pull-on-demand) + parent = self.root_link_vel_w + if self._root_link_lin_vel_w_ta is None: + self._root_link_lin_vel_w_ta = ProxyArray(self._get_lin_vel_from_spatial_vector(parent.warp)) + return self._root_link_lin_vel_w_ta @property - def root_link_ang_vel_w(self) -> wp.array: + def root_link_ang_vel_w(self) -> ProxyArray: """Root link angular velocity in simulation world frame. Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). This quantity is the angular velocity of the actor frame of the root rigid body relative to the world. """ - return self._get_ang_vel_from_spatial_vector(self.root_link_vel_w) + # Access parent property to trigger its getter call (PhysX is pull-on-demand) + parent = self.root_link_vel_w + if self._root_link_ang_vel_w_ta is None: + self._root_link_ang_vel_w_ta = ProxyArray(self._get_ang_vel_from_spatial_vector(parent.warp)) + return self._root_link_ang_vel_w_ta @property - def root_com_pos_w(self) -> wp.array: + def root_com_pos_w(self) -> ProxyArray: """Root center of mass position in simulation world frame. Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). This quantity is the position of the center of mass frame of the root rigid body relative to the world. """ - return self._get_pos_from_transform(self.root_com_pose_w) + # Access parent property to trigger its getter call (PhysX is pull-on-demand) + parent = self.root_com_pose_w + if self._root_com_pos_w_ta is None: + self._root_com_pos_w_ta = ProxyArray(self._get_pos_from_transform(parent.warp)) + return self._root_com_pos_w_ta @property - def root_com_quat_w(self) -> wp.array: + def root_com_quat_w(self) -> ProxyArray: """Root center of mass orientation (x, y, z, w) in simulation world frame. Shape is (num_instances,), dtype = wp.quatf. In torch this resolves to (num_instances, 4). This quantity is the orientation of the principal axes of inertia of the root rigid body relative to the world. """ - return self._get_quat_from_transform(self.root_com_pose_w) + # Access parent property to trigger its getter call (PhysX is pull-on-demand) + parent = self.root_com_pose_w + if self._root_com_quat_w_ta is None: + self._root_com_quat_w_ta = ProxyArray(self._get_quat_from_transform(parent.warp)) + return self._root_com_quat_w_ta @property - def root_com_lin_vel_w(self) -> wp.array: + def root_com_lin_vel_w(self) -> ProxyArray: """Root center of mass linear velocity in simulation world frame. Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). This quantity is the linear velocity of the root rigid body's center of mass frame relative to the world. """ - return self._get_lin_vel_from_spatial_vector(self.root_com_vel_w) + # Access parent property to trigger its getter call (PhysX is pull-on-demand) + parent = self.root_com_vel_w + if self._root_com_lin_vel_w_ta is None: + self._root_com_lin_vel_w_ta = ProxyArray(self._get_lin_vel_from_spatial_vector(parent.warp)) + return self._root_com_lin_vel_w_ta @property - def root_com_ang_vel_w(self) -> wp.array: + def root_com_ang_vel_w(self) -> ProxyArray: """Root center of mass angular velocity in simulation world frame. Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). This quantity is the angular velocity of the root rigid body's center of mass frame relative to the world. """ - return self._get_ang_vel_from_spatial_vector(self.root_com_vel_w) + # Access parent property to trigger its getter call (PhysX is pull-on-demand) + parent = self.root_com_vel_w + if self._root_com_ang_vel_w_ta is None: + self._root_com_ang_vel_w_ta = ProxyArray(self._get_ang_vel_from_spatial_vector(parent.warp)) + return self._root_com_ang_vel_w_ta @property - def body_link_pos_w(self) -> wp.array: + def body_link_pos_w(self) -> ProxyArray: """Positions of all bodies in simulation world frame. Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to (num_instances, num_bodies, 3). This quantity is the position of the articulation bodies' actor frame relative to the world. """ - return self._get_pos_from_transform(self.body_link_pose_w) + # Access parent property to trigger its getter call (PhysX is pull-on-demand) + parent = self.body_link_pose_w + if self._body_link_pos_w_ta is None: + self._body_link_pos_w_ta = ProxyArray(self._get_pos_from_transform(parent.warp)) + return self._body_link_pos_w_ta @property - def body_link_quat_w(self) -> wp.array: + def body_link_quat_w(self) -> ProxyArray: """Orientation (x, y, z, w) of all bodies in simulation world frame. Shape is (num_instances, num_bodies), dtype = wp.quatf. In torch this resolves to (num_instances, num_bodies, 4). This quantity is the orientation of the articulation bodies' actor frame relative to the world. """ - return self._get_quat_from_transform(self.body_link_pose_w) + # Access parent property to trigger its getter call (PhysX is pull-on-demand) + parent = self.body_link_pose_w + if self._body_link_quat_w_ta is None: + self._body_link_quat_w_ta = ProxyArray(self._get_quat_from_transform(parent.warp)) + return self._body_link_quat_w_ta @property - def body_link_lin_vel_w(self) -> wp.array: + def body_link_lin_vel_w(self) -> ProxyArray: """Linear velocity of all bodies in simulation world frame. Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to (num_instances, num_bodies, 3). This quantity is the linear velocity of the articulation bodies' actor frame relative to the world. """ - return self._get_lin_vel_from_spatial_vector(self.body_link_vel_w) + # Access parent property to trigger its getter call (PhysX is pull-on-demand) + parent = self.body_link_vel_w + if self._body_link_lin_vel_w_ta is None: + self._body_link_lin_vel_w_ta = ProxyArray(self._get_lin_vel_from_spatial_vector(parent.warp)) + return self._body_link_lin_vel_w_ta @property - def body_link_ang_vel_w(self) -> wp.array: + def body_link_ang_vel_w(self) -> ProxyArray: """Angular velocity of all bodies in simulation world frame. Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to (num_instances, num_bodies, 3). This quantity is the angular velocity of the articulation bodies' actor frame relative to the world. """ - return self._get_ang_vel_from_spatial_vector(self.body_link_vel_w) + # Access parent property to trigger its getter call (PhysX is pull-on-demand) + parent = self.body_link_vel_w + if self._body_link_ang_vel_w_ta is None: + self._body_link_ang_vel_w_ta = ProxyArray(self._get_ang_vel_from_spatial_vector(parent.warp)) + return self._body_link_ang_vel_w_ta @property - def body_com_pos_w(self) -> wp.array: + def body_com_pos_w(self) -> ProxyArray: """Positions of all bodies' center of mass in simulation world frame. Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to (num_instances, num_bodies, 3). This quantity is the position of the articulation bodies' center of mass frame. """ - return self._get_pos_from_transform(self.body_com_pose_w) + # Access parent property to trigger its getter call (PhysX is pull-on-demand) + parent = self.body_com_pose_w + if self._body_com_pos_w_ta is None: + self._body_com_pos_w_ta = ProxyArray(self._get_pos_from_transform(parent.warp)) + return self._body_com_pos_w_ta @property - def body_com_quat_w(self) -> wp.array: + def body_com_quat_w(self) -> ProxyArray: """Orientation (x, y, z, w) of the principal axes of inertia of all bodies in simulation world frame. Shape is (num_instances, num_bodies), dtype = wp.quatf. In torch this resolves to (num_instances, num_bodies, 4). This quantity is the orientation of the articulation bodies' principal axes of inertia. """ - return self._get_quat_from_transform(self.body_com_pose_w) + # Access parent property to trigger its getter call (PhysX is pull-on-demand) + parent = self.body_com_pose_w + if self._body_com_quat_w_ta is None: + self._body_com_quat_w_ta = ProxyArray(self._get_quat_from_transform(parent.warp)) + return self._body_com_quat_w_ta @property - def body_com_lin_vel_w(self) -> wp.array: + def body_com_lin_vel_w(self) -> ProxyArray: """Linear velocity of all bodies in simulation world frame. Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to (num_instances, num_bodies, 3). This quantity is the linear velocity of the articulation bodies' center of mass frame. """ - return self._get_lin_vel_from_spatial_vector(self.body_com_vel_w) + # Access parent property to trigger its getter call (PhysX is pull-on-demand) + parent = self.body_com_vel_w + if self._body_com_lin_vel_w_ta is None: + self._body_com_lin_vel_w_ta = ProxyArray(self._get_lin_vel_from_spatial_vector(parent.warp)) + return self._body_com_lin_vel_w_ta @property - def body_com_ang_vel_w(self) -> wp.array: + def body_com_ang_vel_w(self) -> ProxyArray: """Angular velocity of all bodies in simulation world frame. Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to (num_instances, num_bodies, 3). This quantity is the angular velocity of the articulation bodies' center of mass frame. """ - return self._get_ang_vel_from_spatial_vector(self.body_com_vel_w) + # Access parent property to trigger its getter call (PhysX is pull-on-demand) + parent = self.body_com_vel_w + if self._body_com_ang_vel_w_ta is None: + self._body_com_ang_vel_w_ta = ProxyArray(self._get_ang_vel_from_spatial_vector(parent.warp)) + return self._body_com_ang_vel_w_ta @property - def body_com_lin_acc_w(self) -> wp.array: + def body_com_lin_acc_w(self) -> ProxyArray: """Linear acceleration of all bodies in simulation world frame. Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to (num_instances, num_bodies, 3). This quantity is the linear acceleration of the articulation bodies' center of mass frame. """ - return self._get_lin_vel_from_spatial_vector(self.body_com_acc_w) + # Access parent property to trigger its getter call (PhysX is pull-on-demand) + parent = self.body_com_acc_w + if self._body_com_lin_acc_w_ta is None: + self._body_com_lin_acc_w_ta = ProxyArray(self._get_lin_vel_from_spatial_vector(parent.warp)) + return self._body_com_lin_acc_w_ta @property - def body_com_ang_acc_w(self) -> wp.array: + def body_com_ang_acc_w(self) -> ProxyArray: """Angular acceleration of all bodies in simulation world frame. Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to (num_instances, num_bodies, 3). This quantity is the angular acceleration of the articulation bodies' center of mass frame. """ - return self._get_ang_vel_from_spatial_vector(self.body_com_acc_w) + # Access parent property to trigger its getter call (PhysX is pull-on-demand) + parent = self.body_com_acc_w + if self._body_com_ang_acc_w_ta is None: + self._body_com_ang_acc_w_ta = ProxyArray(self._get_ang_vel_from_spatial_vector(parent.warp)) + return self._body_com_ang_acc_w_ta @property - def body_com_pos_b(self) -> wp.array: + def body_com_pos_b(self) -> ProxyArray: """Center of mass position of all of the bodies in their respective link frames. Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to (num_instances, num_bodies, 3). This quantity is the center of mass location relative to its body's link frame. """ - return self._get_pos_from_transform(self.body_com_pose_b) + # Access parent property to trigger its getter call (PhysX is pull-on-demand) + parent = self.body_com_pose_b + if self._body_com_pos_b_ta is None: + self._body_com_pos_b_ta = ProxyArray(self._get_pos_from_transform(parent.warp)) + return self._body_com_pos_b_ta @property - def body_com_quat_b(self) -> wp.array: + def body_com_quat_b(self) -> ProxyArray: """Orientation (x, y, z, w) of the principal axes of inertia of all of the bodies in their respective link frames. Shape is (num_instances, num_bodies), dtype = wp.quatf. In torch this resolves to @@ -1133,7 +1331,11 @@ def body_com_quat_b(self) -> wp.array: This quantity is the orientation of the principal axes of inertia relative to its body's link frame. """ - return self._get_quat_from_transform(self.body_com_pose_b) + # Access parent property to trigger its getter call (PhysX is pull-on-demand) + parent = self.body_com_pose_b + if self._body_com_quat_b_ta is None: + self._body_com_quat_b_ta = ProxyArray(self._get_quat_from_transform(parent.warp)) + return self._body_com_quat_b_ta def _create_buffers(self) -> None: super()._create_buffers() @@ -1292,6 +1494,116 @@ def _create_buffers(self) -> None: self._body_inertia = wp.clone(self._root_view.get_inertias(), device=self.device) self._default_root_state = None + # Initialize ProxyArray wrappers + self._pin_proxy_arrays() + + def _pin_proxy_arrays(self) -> None: + """Create pinned ProxyArray wrappers for all data buffers. + + This is called once from :meth:`_create_buffers` during initialization. + PhysX tensor API buffers have stable GPU pointers across simulation steps, + so no rebinding is needed (unlike Newton). + """ + # -- Pinned ProxyArray cache (one per read property, lazily created on first access) + # Defaults + self._default_root_pose_ta: ProxyArray | None = None + self._default_root_vel_ta: ProxyArray | None = None + self._default_joint_pos_ta: ProxyArray | None = None + self._default_joint_vel_ta: ProxyArray | None = None + # Joint commands (set into simulation) + self._joint_pos_target_ta: ProxyArray | None = None + self._joint_vel_target_ta: ProxyArray | None = None + self._joint_effort_target_ta: ProxyArray | None = None + # Joint commands (explicit actuator model) + self._computed_torque_ta: ProxyArray | None = None + self._applied_torque_ta: ProxyArray | None = None + # Joint properties + self._joint_stiffness_ta: ProxyArray | None = None + self._joint_damping_ta: ProxyArray | None = None + self._joint_armature_ta: ProxyArray | None = None + self._joint_friction_coeff_ta: ProxyArray | None = None + self._joint_dynamic_friction_coeff_ta: ProxyArray | None = None + self._joint_viscous_friction_coeff_ta: ProxyArray | None = None + self._joint_pos_limits_ta: ProxyArray | None = None + self._joint_vel_limits_ta: ProxyArray | None = None + self._joint_effort_limits_ta: ProxyArray | None = None + # Joint properties (custom) + self._soft_joint_pos_limits_ta: ProxyArray | None = None + self._soft_joint_vel_limits_ta: ProxyArray | None = None + self._gear_ratio_ta: ProxyArray | None = None + # Fixed tendon properties + self._fixed_tendon_stiffness_ta: ProxyArray | None = None + self._fixed_tendon_damping_ta: ProxyArray | None = None + self._fixed_tendon_limit_stiffness_ta: ProxyArray | None = None + self._fixed_tendon_rest_length_ta: ProxyArray | None = None + self._fixed_tendon_offset_ta: ProxyArray | None = None + self._fixed_tendon_pos_limits_ta: ProxyArray | None = None + # Spatial tendon properties + self._spatial_tendon_stiffness_ta: ProxyArray | None = None + self._spatial_tendon_damping_ta: ProxyArray | None = None + self._spatial_tendon_limit_stiffness_ta: ProxyArray | None = None + self._spatial_tendon_offset_ta: ProxyArray | None = None + # Root state (timestamped) + self._root_link_pose_w_ta: ProxyArray | None = None + self._root_link_vel_w_ta: ProxyArray | None = None + self._root_com_pose_w_ta: ProxyArray | None = None + self._root_com_vel_w_ta: ProxyArray | None = None + # Body state (timestamped) + self._body_link_pose_w_ta: ProxyArray | None = None + self._body_link_vel_w_ta: ProxyArray | None = None + self._body_com_pose_w_ta: ProxyArray | None = None + self._body_com_vel_w_ta: ProxyArray | None = None + self._body_com_acc_w_ta: ProxyArray | None = None + self._body_com_pose_b_ta: ProxyArray | None = None + self._body_incoming_joint_wrench_b_ta: ProxyArray | None = None + # Body properties + self._body_mass_ta: ProxyArray | None = None + self._body_inertia_ta: ProxyArray | None = None + # Joint state (timestamped) + self._joint_pos_ta: ProxyArray | None = None + self._joint_vel_ta: ProxyArray | None = None + self._joint_acc_ta: ProxyArray | None = None + # Derived properties (timestamped) + self._projected_gravity_b_ta: ProxyArray | None = None + self._heading_w_ta: ProxyArray | None = None + self._root_link_lin_vel_b_ta: ProxyArray | None = None + self._root_link_ang_vel_b_ta: ProxyArray | None = None + self._root_com_lin_vel_b_ta: ProxyArray | None = None + self._root_com_ang_vel_b_ta: ProxyArray | None = None + # Sliced properties (root link) + self._root_link_pos_w_ta: ProxyArray | None = None + self._root_link_quat_w_ta: ProxyArray | None = None + self._root_link_lin_vel_w_ta: ProxyArray | None = None + self._root_link_ang_vel_w_ta: ProxyArray | None = None + # Sliced properties (root com) + self._root_com_pos_w_ta: ProxyArray | None = None + self._root_com_quat_w_ta: ProxyArray | None = None + self._root_com_lin_vel_w_ta: ProxyArray | None = None + self._root_com_ang_vel_w_ta: ProxyArray | None = None + # Sliced properties (body link) + self._body_link_pos_w_ta: ProxyArray | None = None + self._body_link_quat_w_ta: ProxyArray | None = None + self._body_link_lin_vel_w_ta: ProxyArray | None = None + self._body_link_ang_vel_w_ta: ProxyArray | None = None + # Sliced properties (body com) + self._body_com_pos_w_ta: ProxyArray | None = None + self._body_com_quat_w_ta: ProxyArray | None = None + self._body_com_lin_vel_w_ta: ProxyArray | None = None + self._body_com_ang_vel_w_ta: ProxyArray | None = None + self._body_com_lin_acc_w_ta: ProxyArray | None = None + self._body_com_ang_acc_w_ta: ProxyArray | None = None + # Sliced properties (body com in body frame) + self._body_com_pos_b_ta: ProxyArray | None = None + self._body_com_quat_b_ta: ProxyArray | None = None + # Deprecated state-concat properties + self._default_root_state_ta: ProxyArray | None = None + self._root_state_w_ta: ProxyArray | None = None + self._root_link_state_w_ta: ProxyArray | None = None + self._root_com_state_w_ta: ProxyArray | None = None + self._body_state_w_ta: ProxyArray | None = None + self._body_link_state_w_ta: ProxyArray | None = None + self._body_com_state_w_ta: ProxyArray | None = None + """ Internal helpers. """ @@ -1369,7 +1681,7 @@ def _get_ang_vel_from_spatial_vector(self, spatial_vector: wp.array) -> wp.array """ @property - def default_root_state(self) -> wp.array: + def default_root_state(self) -> ProxyArray: """Default root state ``[pos, quat, lin_vel, ang_vel]`` in the local environment frame. @@ -1398,10 +1710,12 @@ def default_root_state(self) -> wp.array: ], device=self.device, ) - return self._default_root_state + if self._default_root_state_ta is None: + self._default_root_state_ta = ProxyArray(self._default_root_state) + return self._default_root_state_ta @property - def root_state_w(self) -> wp.array: + def root_state_w(self) -> ProxyArray: """Deprecated, same as :attr:`root_link_pose_w` and :attr:`root_com_vel_w`.""" warnings.warn( "The `root_state_w` property will be deprecated in a IsaacLab 4.0. Please use `root_link_pose_w` and " @@ -1424,10 +1738,12 @@ def root_state_w(self) -> wp.array: ) self._root_state_w.timestamp = self._sim_timestamp - return self._root_state_w.data + if self._root_state_w_ta is None: + self._root_state_w_ta = ProxyArray(self._root_state_w.data) + return self._root_state_w_ta @property - def root_link_state_w(self) -> wp.array: + def root_link_state_w(self) -> ProxyArray: """Deprecated, same as :attr:`root_link_pose_w` and :attr:`root_link_vel_w`.""" warnings.warn( "The `root_link_state_w` property will be deprecated in a IsaacLab 4.0. Please use `root_link_pose_w` and " @@ -1450,10 +1766,12 @@ def root_link_state_w(self) -> wp.array: ) self._root_link_state_w.timestamp = self._sim_timestamp - return self._root_link_state_w.data + if self._root_link_state_w_ta is None: + self._root_link_state_w_ta = ProxyArray(self._root_link_state_w.data) + return self._root_link_state_w_ta @property - def root_com_state_w(self) -> wp.array: + def root_com_state_w(self) -> ProxyArray: """Deprecated, same as :attr:`root_com_pose_w` and :attr:`root_com_vel_w`.""" warnings.warn( "The `root_com_state_w` property will be deprecated in a IsaacLab 4.0. Please use `root_com_pose_w` and " @@ -1476,10 +1794,12 @@ def root_com_state_w(self) -> wp.array: ) self._root_com_state_w.timestamp = self._sim_timestamp - return self._root_com_state_w.data + if self._root_com_state_w_ta is None: + self._root_com_state_w_ta = ProxyArray(self._root_com_state_w.data) + return self._root_com_state_w_ta @property - def body_state_w(self): + def body_state_w(self) -> ProxyArray: """State of all bodies `[pos, quat, lin_vel, ang_vel]` in simulation world frame. Shape is (num_instances, num_bodies, 13). @@ -1507,10 +1827,12 @@ def body_state_w(self): ) self._body_state_w.timestamp = self._sim_timestamp - return self._body_state_w.data + if self._body_state_w_ta is None: + self._body_state_w_ta = ProxyArray(self._body_state_w.data) + return self._body_state_w_ta @property - def body_link_state_w(self): + def body_link_state_w(self) -> ProxyArray: """State of all bodies' link frame`[pos, quat, lin_vel, ang_vel]` in simulation world frame. Shape is (num_instances, num_bodies, 13). @@ -1537,10 +1859,12 @@ def body_link_state_w(self): ) self._body_link_state_w.timestamp = self._sim_timestamp - return self._body_link_state_w.data + if self._body_link_state_w_ta is None: + self._body_link_state_w_ta = ProxyArray(self._body_link_state_w.data) + return self._body_link_state_w_ta @property - def body_com_state_w(self): + def body_com_state_w(self) -> ProxyArray: """State of all bodies center of mass `[pos, quat, lin_vel, ang_vel]` in simulation world frame. Shape is (num_instances, num_bodies, 13). @@ -1569,4 +1893,6 @@ def body_com_state_w(self): ) self._body_com_state_w.timestamp = self._sim_timestamp - return self._body_com_state_w.data + if self._body_com_state_w_ta is None: + self._body_com_state_w_ta = ProxyArray(self._body_com_state_w.data) + return self._body_com_state_w_ta diff --git a/source/isaaclab_physx/isaaclab_physx/assets/deformable_object/deformable_object.py b/source/isaaclab_physx/isaaclab_physx/assets/deformable_object/deformable_object.py index 3bc6ed2b8b25..d6476ba24a27 100644 --- a/source/isaaclab_physx/isaaclab_physx/assets/deformable_object/deformable_object.py +++ b/source/isaaclab_physx/isaaclab_physx/assets/deformable_object/deformable_object.py @@ -21,6 +21,7 @@ import isaaclab.utils.math as math_utils from isaaclab.assets.asset_base import AssetBase from isaaclab.markers import VisualizationMarkers +from isaaclab.utils.warp import ProxyArray from isaaclab_physx.physics import PhysxManager as SimulationManager @@ -407,7 +408,7 @@ def write_nodal_kinematic_target_to_sim_index( ) # set into simulation self.root_view.set_simulation_nodal_kinematic_targets( - self._data.nodal_kinematic_target.view(wp.float32), indices=env_ids + self._data.nodal_kinematic_target.warp.view(wp.float32), indices=env_ids ) def write_nodal_kinematic_target_to_sim_mask( @@ -700,35 +701,35 @@ def _create_buffers(self): (self.num_instances, self.max_sim_vertices_per_body), dtype=wp.vec3f, device=self.device ) # compute default nodal state as vec6f - self._data.default_nodal_state_w = wp.zeros( + default_nodal_state = wp.zeros( (self.num_instances, self.max_sim_vertices_per_body), dtype=vec6f, device=self.device ) wp.launch( compute_nodal_state_w, dim=(self.num_instances, self.max_sim_vertices_per_body), inputs=[nodal_positions, nodal_velocities], - outputs=[self._data.default_nodal_state_w], + outputs=[default_nodal_state], device=self.device, ) + self._data.default_nodal_state_w = ProxyArray(default_nodal_state) # kinematic targets (only for volume deformables, surface deformables do not support kinematic targets) if self._deformable_type == "volume": # kinematic targets — allocate our own buffer and copy from PhysX kinematic_raw = self.root_view.get_simulation_nodal_kinematic_targets() # (N, V, 4) float32 kinematic_view = kinematic_raw.view(wp.vec4f).reshape((self.num_instances, self.max_sim_vertices_per_body)) - self._data.nodal_kinematic_target = wp.zeros( + kinematic_target = wp.zeros( (self.num_instances, self.max_sim_vertices_per_body), dtype=wp.vec4f, device=self.device ) - wp.copy(self._data.nodal_kinematic_target, kinematic_view) + wp.copy(kinematic_target, kinematic_view) # set all nodes as non-kinematic targets by default (flag = 1.0) wp.launch( set_kinematic_flags_to_one, dim=(self.num_instances * self.max_sim_vertices_per_body,), - inputs=[ - self._data.nodal_kinematic_target.reshape((self.num_instances * self.max_sim_vertices_per_body,)) - ], + inputs=[kinematic_target.reshape((self.num_instances * self.max_sim_vertices_per_body,))], device=self.device, ) + self._data.nodal_kinematic_target = ProxyArray(kinematic_target) """ Internal simulation callbacks. @@ -750,7 +751,7 @@ def _debug_vis_callback(self, event): # check where to visualize, kinematic targets only supported for volume deformables num_enabled = 0 if self._deformable_type == "volume": - kinematic_target_torch = wp.to_torch(self.data.nodal_kinematic_target) + kinematic_target_torch = self.data.nodal_kinematic_target.torch targets_enabled = kinematic_target_torch[:, :, 3] == 0.0 num_enabled = int(torch.sum(targets_enabled).item()) # get positions if any targets are enabled diff --git a/source/isaaclab_physx/isaaclab_physx/assets/deformable_object/deformable_object_data.py b/source/isaaclab_physx/isaaclab_physx/assets/deformable_object/deformable_object_data.py index ccac2cf05b71..7f5d604205b5 100644 --- a/source/isaaclab_physx/isaaclab_physx/assets/deformable_object/deformable_object_data.py +++ b/source/isaaclab_physx/isaaclab_physx/assets/deformable_object/deformable_object_data.py @@ -10,6 +10,7 @@ import omni.physics.tensors.api as physx from isaaclab.utils.buffers import TimestampedBufferWarp as TimestampedBuffer +from isaaclab.utils.warp import ProxyArray from .kernels import compute_mean_vec3f_over_vertices, compute_nodal_state_w, vec6f @@ -66,6 +67,13 @@ def __init__(self, root_view: physx.DeformableBodyView, device: str): self._root_pos_w = TimestampedBuffer((self._num_instances,), device, wp.vec3f) self._root_vel_w = TimestampedBuffer((self._num_instances,), device, wp.vec3f) + # -- Pinned ProxyArray cache (one per read property, lazily created on first access) + self._nodal_pos_w_ta: ProxyArray | None = None + self._nodal_vel_w_ta: ProxyArray | None = None + self._nodal_state_w_ta: ProxyArray | None = None + self._root_pos_w_ta: ProxyArray | None = None + self._root_vel_w_ta: ProxyArray | None = None + def update(self, dt: float): """Updates the data for the deformable object. @@ -79,7 +87,7 @@ def update(self, dt: float): # Defaults. ## - default_nodal_state_w: wp.array = None + default_nodal_state_w: ProxyArray = None """Default nodal state ``[nodal_pos, nodal_vel]`` in simulation world frame. Shape is (num_instances, max_sim_vertices_per_body) with dtype vec6f. """ @@ -88,7 +96,7 @@ def update(self, dt: float): # Kinematic commands ## - nodal_kinematic_target: wp.array = None + nodal_kinematic_target: ProxyArray = None """Simulation mesh kinematic targets for the deformable bodies. Shape is (num_instances, max_sim_vertices_per_body) with dtype vec4f. @@ -103,7 +111,7 @@ def update(self, dt: float): ## @property - def nodal_pos_w(self) -> wp.array: + def nodal_pos_w(self) -> ProxyArray: """Nodal positions in simulation world frame. Shape is (num_instances, max_sim_vertices_per_body) vec3f.""" if self._nodal_pos_w.timestamp < self._sim_timestamp: # get_simulation_nodal_positions() returns (N, V, 3) float32 — view as (N, V) vec3f @@ -113,10 +121,15 @@ def nodal_pos_w(self) -> wp.array: .reshape((self._num_instances, self._max_sim_vertices)) ) self._nodal_pos_w.timestamp = self._sim_timestamp - return self._nodal_pos_w.data + # Rebind ProxyArray since .data was replaced with a new wp.array + if self._nodal_pos_w_ta is not None: + self._nodal_pos_w_ta = ProxyArray(self._nodal_pos_w.data) + if self._nodal_pos_w_ta is None: + self._nodal_pos_w_ta = ProxyArray(self._nodal_pos_w.data) + return self._nodal_pos_w_ta @property - def nodal_vel_w(self) -> wp.array: + def nodal_vel_w(self) -> ProxyArray: """Nodal velocities in simulation world frame. Shape is (num_instances, max_sim_vertices_per_body) vec3f.""" if self._nodal_vel_w.timestamp < self._sim_timestamp: self._nodal_vel_w.data = ( @@ -125,10 +138,15 @@ def nodal_vel_w(self) -> wp.array: .reshape((self._num_instances, self._max_sim_vertices)) ) self._nodal_vel_w.timestamp = self._sim_timestamp - return self._nodal_vel_w.data + # Rebind ProxyArray since .data was replaced with a new wp.array + if self._nodal_vel_w_ta is not None: + self._nodal_vel_w_ta = ProxyArray(self._nodal_vel_w.data) + if self._nodal_vel_w_ta is None: + self._nodal_vel_w_ta = ProxyArray(self._nodal_vel_w.data) + return self._nodal_vel_w_ta @property - def nodal_state_w(self) -> wp.array: + def nodal_state_w(self) -> ProxyArray: """Nodal state ``[nodal_pos, nodal_vel]`` in simulation world frame. Shape is (num_instances, max_sim_vertices_per_body) vec6f. """ @@ -141,14 +159,16 @@ def nodal_state_w(self) -> wp.array: device=self.device, ) self._nodal_state_w.timestamp = self._sim_timestamp - return self._nodal_state_w.data + if self._nodal_state_w_ta is None: + self._nodal_state_w_ta = ProxyArray(self._nodal_state_w.data) + return self._nodal_state_w_ta ## # Derived properties. ## @property - def root_pos_w(self) -> wp.array: + def root_pos_w(self) -> ProxyArray: """Root position from nodal positions of the simulation mesh for the deformable bodies in simulation world frame. Shape is (num_instances, 3). @@ -163,10 +183,12 @@ def root_pos_w(self) -> wp.array: device=self.device, ) self._root_pos_w.timestamp = self._sim_timestamp - return self._root_pos_w.data + if self._root_pos_w_ta is None: + self._root_pos_w_ta = ProxyArray(self._root_pos_w.data) + return self._root_pos_w_ta @property - def root_vel_w(self) -> wp.array: + def root_vel_w(self) -> ProxyArray: """Root velocity from vertex velocities for the deformable bodies in simulation world frame. Shape is (num_instances, 3). @@ -181,4 +203,6 @@ def root_vel_w(self) -> wp.array: device=self.device, ) self._root_vel_w.timestamp = self._sim_timestamp - return self._root_vel_w.data + if self._root_vel_w_ta is None: + self._root_vel_w_ta = ProxyArray(self._root_vel_w.data) + return self._root_vel_w_ta 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 25abb8fd22a3..6c0be5d2a6a1 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 @@ -156,8 +156,8 @@ def write_data_to_sim(self) -> None: composer = self._permanent_wrench_composer composer.compose_to_body_frame() self.root_view.apply_forces_and_torques_at_position( - force_data=composer.out_force_b.flatten().view(wp.float32), - torque_data=composer.out_torque_b.flatten().view(wp.float32), + force_data=composer.out_force_b.warp.flatten().view(wp.float32), + torque_data=composer.out_torque_b.warp.flatten().view(wp.float32), position_data=None, indices=self._ALL_INDICES, is_global=False, 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 2bd71b799b7e..a8e674dc6bf8 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 @@ -15,6 +15,7 @@ from isaaclab.assets.rigid_object.base_rigid_object_data import BaseRigidObjectData from isaaclab.utils.buffers import TimestampedBufferWarp as TimestampedBuffer from isaaclab.utils.math import normalize +from isaaclab.utils.warp import ProxyArray from isaaclab_physx.assets import kernels as shared_kernels from isaaclab_physx.physics import PhysxManager as SimulationManager @@ -45,6 +46,16 @@ class RigidObjectData(BaseRigidObjectData): The data is lazily updated, meaning that the data is only updated when it is accessed. This is useful when the data is expensive to compute or retrieve. The data is updated when the timestamp of the buffer is older than the current simulation timestamp. The timestamp is updated whenever the data is updated. + + .. note:: + **Pull-to-refresh model.** Properties pull fresh data from the PhysX tensor API on first + access per timestamp and cache the result. This differs from Newton, where buffers are + refreshed automatically by the simulation. + + .. note:: + **ProxyArray pointer stability.** Each :class:`ProxyArray` wrapper is created once and + reused because the PhysX tensor API returns views into stable, pre-allocated GPU buffers + whose device pointer does not change across simulation steps. """ __backend_name__: str = "physx" @@ -78,8 +89,8 @@ def __init__(self, root_view: physx.RigidBodyView, device: str): forward_vec = torch.tensor((1.0, 0.0, 0.0), device=self.device).repeat(self._root_view.count, 1) # Initialize constants - self.GRAVITY_VEC_W = wp.from_torch(gravity_dir, dtype=wp.vec3f) - self.FORWARD_VEC_B = wp.from_torch(forward_vec, dtype=wp.vec3f) + self.GRAVITY_VEC_W = ProxyArray(wp.from_torch(gravity_dir, dtype=wp.vec3f)) + self.FORWARD_VEC_B = ProxyArray(wp.from_torch(forward_vec, dtype=wp.vec3f)) self._create_buffers() @@ -126,13 +137,15 @@ def update(self, dt: float) -> None: """ @property - def default_root_pose(self) -> wp.array: + def default_root_pose(self) -> ProxyArray: """Default root pose ``[pos, quat]`` in local environment frame. Shape is (num_instances,), dtype = wp.transformf. In torch this resolves to (num_instances, 7). The position and quaternion are of the rigid body's actor frame. """ - return self._default_root_pose + if self._default_root_pose_ta is None: + self._default_root_pose_ta = ProxyArray(self._default_root_pose) + return self._default_root_pose_ta @default_root_pose.setter def default_root_pose(self, value: wp.array) -> None: @@ -149,13 +162,15 @@ def default_root_pose(self, value: wp.array) -> None: self._default_root_pose.assign(value) @property - def default_root_vel(self) -> wp.array: + def default_root_vel(self) -> ProxyArray: """Default root velocity ``[lin_vel, ang_vel]`` in local environment frame. Shape is (num_instances,), dtype = wp.spatial_vectorf. In torch this resolves to (num_instances, 6). The linear and angular velocities are of the rigid body's center of mass frame. """ - return self._default_root_vel + if self._default_root_vel_ta is None: + self._default_root_vel_ta = ProxyArray(self._default_root_vel) + return self._default_root_vel_ta @default_root_vel.setter def default_root_vel(self, value: wp.array) -> None: @@ -176,7 +191,7 @@ def default_root_vel(self, value: wp.array) -> None: """ @property - def root_link_pose_w(self) -> wp.array: + def root_link_pose_w(self) -> ProxyArray: """Root link pose ``[pos, quat]`` in simulation world frame. Shape is (num_instances,), dtype = wp.transformf. In torch this resolves to (num_instances, 7). @@ -188,10 +203,12 @@ def root_link_pose_w(self) -> wp.array: self._root_link_pose_w.data = self._root_view.get_transforms().view(wp.transformf) self._root_link_pose_w.timestamp = self._sim_timestamp - return self._root_link_pose_w.data + if self._root_link_pose_w_ta is None: + self._root_link_pose_w_ta = ProxyArray(self._root_link_pose_w.data) + return self._root_link_pose_w_ta @property - def root_link_vel_w(self) -> wp.array: + def root_link_vel_w(self) -> ProxyArray: """Root link velocity ``[lin_vel, ang_vel]`` in simulation world frame. Shape is (num_instances,), dtype = wp.spatial_vectorf. In torch this resolves to (num_instances, 6). @@ -215,10 +232,12 @@ def root_link_vel_w(self) -> wp.array: ) self._root_link_vel_w.timestamp = self._sim_timestamp - return self._root_link_vel_w.data + if self._root_link_vel_w_ta is None: + self._root_link_vel_w_ta = ProxyArray(self._root_link_vel_w.data) + return self._root_link_vel_w_ta @property - def root_com_pose_w(self) -> wp.array: + def root_com_pose_w(self) -> ProxyArray: """Root center of mass pose ``[pos, quat]`` in simulation world frame. Shape is (num_instances,), dtype = wp.transformf. In torch this resolves to (num_instances, 7). @@ -241,10 +260,12 @@ def root_com_pose_w(self) -> wp.array: ) self._root_com_pose_w.timestamp = self._sim_timestamp - return self._root_com_pose_w.data + if self._root_com_pose_w_ta is None: + self._root_com_pose_w_ta = ProxyArray(self._root_com_pose_w.data) + return self._root_com_pose_w_ta @property - def root_com_vel_w(self) -> wp.array: + def root_com_vel_w(self) -> ProxyArray: """Root center of mass velocity ``[lin_vel, ang_vel]`` in simulation world frame. Shape is (num_instances,), dtype = wp.spatial_vectorf. In torch this resolves to (num_instances, 6). @@ -255,72 +276,90 @@ def root_com_vel_w(self) -> wp.array: self._root_com_vel_w.data = self._root_view.get_velocities().view(wp.spatial_vectorf) self._root_com_vel_w.timestamp = self._sim_timestamp - return self._root_com_vel_w.data + if self._root_com_vel_w_ta is None: + self._root_com_vel_w_ta = ProxyArray(self._root_com_vel_w.data) + return self._root_com_vel_w_ta """ Body state properties. """ @property - def body_mass(self) -> wp.array: + def body_mass(self) -> ProxyArray: """Mass of all bodies in the simulation world frame. Shape is (num_instances, 1), dtype = wp.float32. In torch this resolves to (num_instances, 1). """ - return self._body_mass + if self._body_mass_ta is None: + self._body_mass_ta = ProxyArray(self._body_mass) + return self._body_mass_ta @property - def body_inertia(self) -> wp.array: + def body_inertia(self) -> ProxyArray: """Inertia of all bodies in the simulation world frame. Shape is (num_instances, 1, 9), dtype = wp.float32. In torch this resolves to (num_instances, 1, 9). """ - return self._body_inertia + if self._body_inertia_ta is None: + self._body_inertia_ta = ProxyArray(self._body_inertia) + return self._body_inertia_ta @property - def body_link_pose_w(self) -> wp.array: + def body_link_pose_w(self) -> ProxyArray: """Body link pose ``[pos, quat]`` in simulation world frame. Shape is (num_instances, 1), dtype = wp.transformf. In torch this resolves to (num_instances, 1, 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. """ - return self.root_link_pose_w.reshape((self._num_instances, 1)) + parent = self.root_link_pose_w + if self._body_link_pose_w_ta is None: + self._body_link_pose_w_ta = ProxyArray(parent.warp.reshape((self._num_instances, 1))) + return self._body_link_pose_w_ta @property - def body_link_vel_w(self) -> wp.array: + def body_link_vel_w(self) -> ProxyArray: """Body link velocity ``[lin_vel, ang_vel]`` in simulation world frame. Shape is (num_instances, 1), dtype = wp.spatial_vectorf. In torch this resolves to (num_instances, 1, 6). This quantity contains the linear and angular velocities of the actor frame of the root rigid body relative to the world. """ - return self.root_link_vel_w.reshape((self._num_instances, 1)) + parent = self.root_link_vel_w + if self._body_link_vel_w_ta is None: + self._body_link_vel_w_ta = ProxyArray(parent.warp.reshape((self._num_instances, 1))) + return self._body_link_vel_w_ta @property - def body_com_pose_w(self) -> wp.array: + def body_com_pose_w(self) -> ProxyArray: """Body center of mass pose ``[pos, quat]`` in simulation world frame. Shape is (num_instances, 1), dtype = wp.transformf. In torch this resolves to (num_instances, 1, 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. """ - return self.root_com_pose_w.reshape((self._num_instances, 1)) + parent = self.root_com_pose_w + if self._body_com_pose_w_ta is None: + self._body_com_pose_w_ta = ProxyArray(parent.warp.reshape((self._num_instances, 1))) + return self._body_com_pose_w_ta @property - def body_com_vel_w(self) -> wp.array: + def body_com_vel_w(self) -> ProxyArray: """Body center of mass velocity ``[lin_vel, ang_vel]`` in simulation world frame. Shape is (num_instances, 1), dtype = wp.spatial_vectorf. In torch this resolves to (num_instances, 1, 6). This quantity contains the linear and angular velocities of the root rigid body's center of mass frame relative to the world. """ - return self.root_com_vel_w.reshape((self._num_instances, 1)) + parent = self.root_com_vel_w + if self._body_com_vel_w_ta is None: + self._body_com_vel_w_ta = ProxyArray(parent.warp.reshape((self._num_instances, 1))) + return self._body_com_vel_w_ta @property - def body_com_acc_w(self) -> wp.array: + def body_com_acc_w(self) -> ProxyArray: """Acceleration of all bodies ``[lin_acc, ang_acc]`` in the simulation world frame. Shape is (num_instances, 1), dtype = wp.spatial_vectorf. In torch this resolves to (num_instances, 1, 6). @@ -332,10 +371,12 @@ def body_com_acc_w(self) -> wp.array: ) self._body_com_acc_w.timestamp = self._sim_timestamp - return self._body_com_acc_w.data + if self._body_com_acc_w_ta is None: + self._body_com_acc_w_ta = ProxyArray(self._body_com_acc_w.data) + return self._body_com_acc_w_ta @property - def body_com_pose_b(self) -> wp.array: + def body_com_pose_b(self) -> ProxyArray: """Center of mass pose ``[pos, quat]`` of all bodies in their respective body's link frames. Shape is (num_instances, 1), dtype = wp.transformf. In torch this resolves to (num_instances, 1, 7). @@ -349,14 +390,16 @@ def body_com_pose_b(self) -> wp.array: ) self._body_com_pose_b.timestamp = self._sim_timestamp - return self._body_com_pose_b.data + if self._body_com_pose_b_ta is None: + self._body_com_pose_b_ta = ProxyArray(self._body_com_pose_b.data) + return self._body_com_pose_b_ta """ Derived Properties. """ @property - def projected_gravity_b(self) -> wp.array: + def projected_gravity_b(self) -> ProxyArray: """Projection of the gravity direction on base frame. Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). @@ -370,10 +413,12 @@ def projected_gravity_b(self) -> wp.array: device=self.device, ) self._projected_gravity_b.timestamp = self._sim_timestamp - return self._projected_gravity_b.data + if self._projected_gravity_b_ta is None: + self._projected_gravity_b_ta = ProxyArray(self._projected_gravity_b.data) + return self._projected_gravity_b_ta @property - def heading_w(self) -> wp.array: + def heading_w(self) -> ProxyArray: """Yaw heading of the base frame (in radians). Shape is (num_instances,), dtype = wp.float32. In torch this resolves to (num_instances,). @@ -391,10 +436,12 @@ def heading_w(self) -> wp.array: device=self.device, ) self._heading_w.timestamp = self._sim_timestamp - return self._heading_w.data + if self._heading_w_ta is None: + self._heading_w_ta = ProxyArray(self._heading_w.data) + return self._heading_w_ta @property - def root_link_lin_vel_b(self) -> wp.array: + def root_link_lin_vel_b(self) -> ProxyArray: """Root link linear velocity in base frame. Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). @@ -410,10 +457,12 @@ def root_link_lin_vel_b(self) -> wp.array: device=self.device, ) self._root_link_lin_vel_b.timestamp = self._sim_timestamp - return self._root_link_lin_vel_b.data + if self._root_link_lin_vel_b_ta is None: + self._root_link_lin_vel_b_ta = ProxyArray(self._root_link_lin_vel_b.data) + return self._root_link_lin_vel_b_ta @property - def root_link_ang_vel_b(self) -> wp.array: + def root_link_ang_vel_b(self) -> ProxyArray: """Root link angular velocity in base frame. Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). @@ -429,10 +478,12 @@ def root_link_ang_vel_b(self) -> wp.array: device=self.device, ) self._root_link_ang_vel_b.timestamp = self._sim_timestamp - return self._root_link_ang_vel_b.data + if self._root_link_ang_vel_b_ta is None: + self._root_link_ang_vel_b_ta = ProxyArray(self._root_link_ang_vel_b.data) + return self._root_link_ang_vel_b_ta @property - def root_com_lin_vel_b(self) -> wp.array: + def root_com_lin_vel_b(self) -> ProxyArray: """Root center of mass linear velocity in base frame. Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). @@ -448,10 +499,12 @@ def root_com_lin_vel_b(self) -> wp.array: device=self.device, ) self._root_com_lin_vel_b.timestamp = self._sim_timestamp - return self._root_com_lin_vel_b.data + if self._root_com_lin_vel_b_ta is None: + self._root_com_lin_vel_b_ta = ProxyArray(self._root_com_lin_vel_b.data) + return self._root_com_lin_vel_b_ta @property - def root_com_ang_vel_b(self) -> wp.array: + def root_com_ang_vel_b(self) -> ProxyArray: """Root center of mass angular velocity in base frame. Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). @@ -467,192 +520,254 @@ def root_com_ang_vel_b(self) -> wp.array: device=self.device, ) self._root_com_ang_vel_b.timestamp = self._sim_timestamp - return self._root_com_ang_vel_b.data + if self._root_com_ang_vel_b_ta is None: + self._root_com_ang_vel_b_ta = ProxyArray(self._root_com_ang_vel_b.data) + return self._root_com_ang_vel_b_ta """ Sliced properties. """ @property - def root_link_pos_w(self) -> wp.array: + def root_link_pos_w(self) -> ProxyArray: """Root link position in simulation world frame. Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). This quantity is the position of the actor frame of the root rigid body relative to the world. """ - return self._get_pos_from_transform(self.root_link_pose_w) + parent = self.root_link_pose_w + if self._root_link_pos_w_ta is None: + self._root_link_pos_w_ta = ProxyArray(self._get_pos_from_transform(parent.warp)) + return self._root_link_pos_w_ta @property - def root_link_quat_w(self) -> wp.array: + def root_link_quat_w(self) -> ProxyArray: """Root link orientation (x, y, z, w) in simulation world frame. Shape is (num_instances,), dtype = wp.quatf. In torch this resolves to (num_instances, 4). This quantity is the orientation of the actor frame of the root rigid body. """ - return self._get_quat_from_transform(self.root_link_pose_w) + parent = self.root_link_pose_w + if self._root_link_quat_w_ta is None: + self._root_link_quat_w_ta = ProxyArray(self._get_quat_from_transform(parent.warp)) + return self._root_link_quat_w_ta @property - def root_link_lin_vel_w(self) -> wp.array: + def root_link_lin_vel_w(self) -> ProxyArray: """Root linear velocity in simulation world frame. Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). This quantity is the linear velocity of the root rigid body's actor frame relative to the world. """ - return self._get_lin_vel_from_spatial_vector(self.root_link_vel_w) + parent = self.root_link_vel_w + if self._root_link_lin_vel_w_ta is None: + self._root_link_lin_vel_w_ta = ProxyArray(self._get_lin_vel_from_spatial_vector(parent.warp)) + return self._root_link_lin_vel_w_ta @property - def root_link_ang_vel_w(self) -> wp.array: + def root_link_ang_vel_w(self) -> ProxyArray: """Root link angular velocity in simulation world frame. Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). This quantity is the angular velocity of the actor frame of the root rigid body relative to the world. """ - return self._get_ang_vel_from_spatial_vector(self.root_link_vel_w) + parent = self.root_link_vel_w + if self._root_link_ang_vel_w_ta is None: + self._root_link_ang_vel_w_ta = ProxyArray(self._get_ang_vel_from_spatial_vector(parent.warp)) + return self._root_link_ang_vel_w_ta @property - def root_com_pos_w(self) -> wp.array: + def root_com_pos_w(self) -> ProxyArray: """Root center of mass position in simulation world frame. Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). This quantity is the position of the center of mass frame of the root rigid body relative to the world. """ - return self._get_pos_from_transform(self.root_com_pose_w) + parent = self.root_com_pose_w + if self._root_com_pos_w_ta is None: + self._root_com_pos_w_ta = ProxyArray(self._get_pos_from_transform(parent.warp)) + return self._root_com_pos_w_ta @property - def root_com_quat_w(self) -> wp.array: + def root_com_quat_w(self) -> ProxyArray: """Root center of mass orientation (x, y, z, w) in simulation world frame. Shape is (num_instances,), dtype = wp.quatf. In torch this resolves to (num_instances, 4). This quantity is the orientation of the principal axes of inertia of the root rigid body relative to the world. """ - return self._get_quat_from_transform(self.root_com_pose_w) + parent = self.root_com_pose_w + if self._root_com_quat_w_ta is None: + self._root_com_quat_w_ta = ProxyArray(self._get_quat_from_transform(parent.warp)) + return self._root_com_quat_w_ta @property - def root_com_lin_vel_w(self) -> wp.array: + def root_com_lin_vel_w(self) -> ProxyArray: """Root center of mass linear velocity in simulation world frame. Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). This quantity is the linear velocity of the root rigid body's center of mass frame relative to the world. """ - return self._get_lin_vel_from_spatial_vector(self.root_com_vel_w) + parent = self.root_com_vel_w + if self._root_com_lin_vel_w_ta is None: + self._root_com_lin_vel_w_ta = ProxyArray(self._get_lin_vel_from_spatial_vector(parent.warp)) + return self._root_com_lin_vel_w_ta @property - def root_com_ang_vel_w(self) -> wp.array: + def root_com_ang_vel_w(self) -> ProxyArray: """Root center of mass angular velocity in simulation world frame. Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). This quantity is the angular velocity of the root rigid body's center of mass frame relative to the world. """ - return self._get_ang_vel_from_spatial_vector(self.root_com_vel_w) + parent = self.root_com_vel_w + if self._root_com_ang_vel_w_ta is None: + self._root_com_ang_vel_w_ta = ProxyArray(self._get_ang_vel_from_spatial_vector(parent.warp)) + return self._root_com_ang_vel_w_ta @property - def body_link_pos_w(self) -> wp.array: + def body_link_pos_w(self) -> ProxyArray: """Positions of all bodies in simulation world frame. Shape is (num_instances, 1), dtype = wp.vec3f. In torch this resolves to (num_instances, 1, 3). This quantity is the position of the rigid bodies' actor frame relative to the world. """ - return self._get_pos_from_transform(self.body_link_pose_w) + parent = self.body_link_pose_w + if self._body_link_pos_w_ta is None: + self._body_link_pos_w_ta = ProxyArray(self._get_pos_from_transform(parent.warp)) + return self._body_link_pos_w_ta @property - def body_link_quat_w(self) -> wp.array: + def body_link_quat_w(self) -> ProxyArray: """Orientation (x, y, z, w) of all bodies in simulation world frame. Shape is (num_instances, 1), dtype = wp.quatf. In torch this resolves to (num_instances, 1, 4). This quantity is the orientation of the rigid bodies' actor frame relative to the world. """ - return self._get_quat_from_transform(self.body_link_pose_w) + parent = self.body_link_pose_w + if self._body_link_quat_w_ta is None: + self._body_link_quat_w_ta = ProxyArray(self._get_quat_from_transform(parent.warp)) + return self._body_link_quat_w_ta @property - def body_link_lin_vel_w(self) -> wp.array: + def body_link_lin_vel_w(self) -> ProxyArray: """Linear velocity of all bodies in simulation world frame. Shape is (num_instances, 1), dtype = wp.vec3f. In torch this resolves to (num_instances, 1, 3). This quantity is the linear velocity of the rigid bodies' actor frame relative to the world. """ - return self._get_lin_vel_from_spatial_vector(self.body_link_vel_w) + parent = self.body_link_vel_w + if self._body_link_lin_vel_w_ta is None: + self._body_link_lin_vel_w_ta = ProxyArray(self._get_lin_vel_from_spatial_vector(parent.warp)) + return self._body_link_lin_vel_w_ta @property - def body_link_ang_vel_w(self) -> wp.array: + def body_link_ang_vel_w(self) -> ProxyArray: """Angular velocity of all bodies in simulation world frame. Shape is (num_instances, 1), dtype = wp.vec3f. In torch this resolves to (num_instances, 1, 3). This quantity is the angular velocity of the rigid bodies' actor frame relative to the world. """ - return self._get_ang_vel_from_spatial_vector(self.body_link_vel_w) + parent = self.body_link_vel_w + if self._body_link_ang_vel_w_ta is None: + self._body_link_ang_vel_w_ta = ProxyArray(self._get_ang_vel_from_spatial_vector(parent.warp)) + return self._body_link_ang_vel_w_ta @property - def body_com_pos_w(self) -> wp.array: + def body_com_pos_w(self) -> ProxyArray: """Positions of all bodies' center of mass in simulation world frame. Shape is (num_instances, 1), dtype = wp.vec3f. In torch this resolves to (num_instances, 1, 3). This quantity is the position of the rigid bodies' center of mass frame. """ - return self._get_pos_from_transform(self.body_com_pose_w) + parent = self.body_com_pose_w + if self._body_com_pos_w_ta is None: + self._body_com_pos_w_ta = ProxyArray(self._get_pos_from_transform(parent.warp)) + return self._body_com_pos_w_ta @property - def body_com_quat_w(self) -> wp.array: + def body_com_quat_w(self) -> ProxyArray: """Orientation (x, y, z, w) of the principal axes of inertia of all bodies in simulation world frame. Shape is (num_instances, 1), dtype = wp.quatf. In torch this resolves to (num_instances, 1, 4). This quantity is the orientation of the principal axes of inertia of the rigid bodies. """ - return self._get_quat_from_transform(self.body_com_pose_w) + parent = self.body_com_pose_w + if self._body_com_quat_w_ta is None: + self._body_com_quat_w_ta = ProxyArray(self._get_quat_from_transform(parent.warp)) + return self._body_com_quat_w_ta @property - def body_com_lin_vel_w(self) -> wp.array: + def body_com_lin_vel_w(self) -> ProxyArray: """Linear velocity of all bodies in simulation world frame. Shape is (num_instances, 1), dtype = wp.vec3f. In torch this resolves to (num_instances, 1, 3). This quantity is the linear velocity of the rigid bodies' center of mass frame. """ - return self._get_lin_vel_from_spatial_vector(self.body_com_vel_w) + parent = self.body_com_vel_w + if self._body_com_lin_vel_w_ta is None: + self._body_com_lin_vel_w_ta = ProxyArray(self._get_lin_vel_from_spatial_vector(parent.warp)) + return self._body_com_lin_vel_w_ta @property - def body_com_ang_vel_w(self) -> wp.array: + def body_com_ang_vel_w(self) -> ProxyArray: """Angular velocity of all bodies in simulation world frame. Shape is (num_instances, 1), dtype = wp.vec3f. In torch this resolves to (num_instances, 1, 3). This quantity is the angular velocity of the rigid bodies' center of mass frame. """ - return self._get_ang_vel_from_spatial_vector(self.body_com_vel_w) + parent = self.body_com_vel_w + if self._body_com_ang_vel_w_ta is None: + self._body_com_ang_vel_w_ta = ProxyArray(self._get_ang_vel_from_spatial_vector(parent.warp)) + return self._body_com_ang_vel_w_ta @property - def body_com_lin_acc_w(self) -> wp.array: + def body_com_lin_acc_w(self) -> ProxyArray: """Linear acceleration of all bodies in simulation world frame. Shape is (num_instances, 1), dtype = wp.vec3f. In torch this resolves to (num_instances, 1, 3). This quantity is the linear acceleration of the rigid bodies' center of mass frame. """ - return self._get_lin_vel_from_spatial_vector(self.body_com_acc_w) + parent = self.body_com_acc_w + if self._body_com_lin_acc_w_ta is None: + self._body_com_lin_acc_w_ta = ProxyArray(self._get_lin_vel_from_spatial_vector(parent.warp)) + return self._body_com_lin_acc_w_ta @property - def body_com_ang_acc_w(self) -> wp.array: + def body_com_ang_acc_w(self) -> ProxyArray: """Angular acceleration of all bodies in simulation world frame. Shape is (num_instances, 1), dtype = wp.vec3f. In torch this resolves to (num_instances, 1, 3). This quantity is the angular acceleration of the rigid bodies' center of mass frame. """ - return self._get_ang_vel_from_spatial_vector(self.body_com_acc_w) + parent = self.body_com_acc_w + if self._body_com_ang_acc_w_ta is None: + self._body_com_ang_acc_w_ta = ProxyArray(self._get_ang_vel_from_spatial_vector(parent.warp)) + return self._body_com_ang_acc_w_ta @property - def body_com_pos_b(self) -> wp.array: + def body_com_pos_b(self) -> ProxyArray: """Center of mass position of all of the bodies in their respective link frames. Shape is (num_instances, 1), dtype = wp.vec3f. In torch this resolves to (num_instances, 1, 3). This quantity is the center of mass location relative to its body's link frame. """ - return self._get_pos_from_transform(self.body_com_pose_b) + parent = self.body_com_pose_b + if self._body_com_pos_b_ta is None: + self._body_com_pos_b_ta = ProxyArray(self._get_pos_from_transform(parent.warp)) + return self._body_com_pos_b_ta @property - def body_com_quat_b(self) -> wp.array: + def body_com_quat_b(self) -> ProxyArray: """Orientation (x, y, z, w) of the principal axes of inertia of all of the bodies in their respective link frames. Shape is (num_instances, 1), dtype = wp.quatf. In torch this resolves to (num_instances, 1, 4). This quantity is the orientation of the principal axes of inertia relative to its body's link frame. """ - return self._get_quat_from_transform(self.body_com_pose_b) + parent = self.body_com_pose_b + if self._body_com_quat_b_ta is None: + self._body_com_quat_b_ta = ProxyArray(self._get_quat_from_transform(parent.warp)) + return self._body_com_quat_b_ta def _create_buffers(self) -> None: super()._create_buffers() @@ -689,11 +804,81 @@ def _create_buffers(self) -> None: (self._num_instances, 1, 9) ) + # Initialize ProxyArray wrappers + self._pin_proxy_arrays() + + def _pin_proxy_arrays(self) -> None: + """Create pinned ProxyArray wrappers for all data buffers. + + This is called once from :meth:`_create_buffers` during initialization. + PhysX tensor API buffers have stable GPU pointers across simulation steps, + so no rebinding is needed (unlike Newton). + """ + # -- Pinned ProxyArray cache (one per read property, lazily created on first access) + # Defaults + self._default_root_pose_ta: ProxyArray | None = None + self._default_root_vel_ta: ProxyArray | None = None + # Root state (timestamped) + self._root_link_pose_w_ta: ProxyArray | None = None + self._root_link_vel_w_ta: ProxyArray | None = None + self._root_com_pose_w_ta: ProxyArray | None = None + self._root_com_vel_w_ta: ProxyArray | None = None + # Body properties + self._body_mass_ta: ProxyArray | None = None + self._body_inertia_ta: ProxyArray | None = None + # Body state (reshaped from root) + self._body_link_pose_w_ta: ProxyArray | None = None + self._body_link_vel_w_ta: ProxyArray | None = None + self._body_com_pose_w_ta: ProxyArray | None = None + self._body_com_vel_w_ta: ProxyArray | None = None + self._body_com_acc_w_ta: ProxyArray | None = None + self._body_com_pose_b_ta: ProxyArray | None = None + # Derived properties (timestamped) + self._projected_gravity_b_ta: ProxyArray | None = None + self._heading_w_ta: ProxyArray | None = None + self._root_link_lin_vel_b_ta: ProxyArray | None = None + self._root_link_ang_vel_b_ta: ProxyArray | None = None + self._root_com_lin_vel_b_ta: ProxyArray | None = None + self._root_com_ang_vel_b_ta: ProxyArray | None = None + # Sliced properties (root link) + self._root_link_pos_w_ta: ProxyArray | None = None + self._root_link_quat_w_ta: ProxyArray | None = None + self._root_link_lin_vel_w_ta: ProxyArray | None = None + self._root_link_ang_vel_w_ta: ProxyArray | None = None + # Sliced properties (root com) + self._root_com_pos_w_ta: ProxyArray | None = None + self._root_com_quat_w_ta: ProxyArray | None = None + self._root_com_lin_vel_w_ta: ProxyArray | None = None + self._root_com_ang_vel_w_ta: ProxyArray | None = None + # Sliced properties (body link) + self._body_link_pos_w_ta: ProxyArray | None = None + self._body_link_quat_w_ta: ProxyArray | None = None + self._body_link_lin_vel_w_ta: ProxyArray | None = None + self._body_link_ang_vel_w_ta: ProxyArray | None = None + # Sliced properties (body com) + self._body_com_pos_w_ta: ProxyArray | None = None + self._body_com_quat_w_ta: ProxyArray | None = None + self._body_com_lin_vel_w_ta: ProxyArray | None = None + self._body_com_ang_vel_w_ta: ProxyArray | None = None + self._body_com_lin_acc_w_ta: ProxyArray | None = None + self._body_com_ang_acc_w_ta: ProxyArray | None = None + # Sliced properties (body com in body frame) + self._body_com_pos_b_ta: ProxyArray | None = None + self._body_com_quat_b_ta: ProxyArray | None = None + # Deprecated state-concat properties + self._default_root_state_ta: ProxyArray | None = None + self._root_state_w_ta: ProxyArray | None = None + self._root_link_state_w_ta: ProxyArray | None = None + self._root_com_state_w_ta: ProxyArray | None = None + self._body_state_w_ta: ProxyArray | None = None + self._body_link_state_w_ta: ProxyArray | None = None + self._body_com_state_w_ta: ProxyArray | None = None + """ Internal helpers. """ - def _get_pos_from_transform(self, transform: wp.array) -> wp.array: + def _get_pos_from_transform(self, transform: wp.array) -> ProxyArray: """Generates a position array from a transform array.""" return wp.array( ptr=transform.ptr, @@ -703,7 +888,7 @@ def _get_pos_from_transform(self, transform: wp.array) -> wp.array: device=self.device, ) - def _get_quat_from_transform(self, transform: wp.array) -> wp.array: + def _get_quat_from_transform(self, transform: wp.array) -> ProxyArray: """Generates a quaternion array from a transform array.""" return wp.array( ptr=transform.ptr + 3 * 4, @@ -713,7 +898,7 @@ def _get_quat_from_transform(self, transform: wp.array) -> wp.array: device=self.device, ) - def _get_lin_vel_from_spatial_vector(self, spatial_vector: wp.array) -> wp.array: + def _get_lin_vel_from_spatial_vector(self, spatial_vector: wp.array) -> ProxyArray: """Generates a linear velocity array from a spatial vector array.""" return wp.array( ptr=spatial_vector.ptr, @@ -723,7 +908,7 @@ def _get_lin_vel_from_spatial_vector(self, spatial_vector: wp.array) -> wp.array device=self.device, ) - def _get_ang_vel_from_spatial_vector(self, spatial_vector: wp.array) -> wp.array: + def _get_ang_vel_from_spatial_vector(self, spatial_vector: wp.array) -> ProxyArray: """Generates an angular velocity array from a spatial vector array.""" return wp.array( ptr=spatial_vector.ptr + 3 * 4, @@ -738,7 +923,7 @@ def _get_ang_vel_from_spatial_vector(self, spatial_vector: wp.array) -> wp.array """ @property - def default_root_state(self) -> wp.array: + def default_root_state(self) -> ProxyArray: """Default root state ``[pos, quat, lin_vel, ang_vel]`` in local environment frame. The position and quaternion are of the rigid body's actor frame. Meanwhile, the linear and angular velocities @@ -764,10 +949,12 @@ def default_root_state(self) -> wp.array: ], device=self.device, ) - return self._default_root_state + if self._default_root_state_ta is None: + self._default_root_state_ta = ProxyArray(self._default_root_state) + return self._default_root_state_ta @property - def root_state_w(self) -> wp.array: + def root_state_w(self) -> ProxyArray: """Deprecated, same as :attr:`root_link_pose_w` and :attr:`root_com_vel_w`.""" warnings.warn( "The `root_state_w` property will be deprecated in IsaacLab 4.0. Please use `root_link_pose_w` and " @@ -790,10 +977,12 @@ def root_state_w(self) -> wp.array: ) self._root_state_w.timestamp = self._sim_timestamp - return self._root_state_w.data + if self._root_state_w_ta is None: + self._root_state_w_ta = ProxyArray(self._root_state_w.data) + return self._root_state_w_ta @property - def root_link_state_w(self) -> wp.array: + def root_link_state_w(self) -> ProxyArray: """Deprecated, same as :attr:`root_link_pose_w` and :attr:`root_link_vel_w`.""" warnings.warn( "The `root_link_state_w` property will be deprecated in IsaacLab 4.0. Please use `root_link_pose_w` and " @@ -816,10 +1005,12 @@ def root_link_state_w(self) -> wp.array: ) self._root_link_state_w.timestamp = self._sim_timestamp - return self._root_link_state_w.data + if self._root_link_state_w_ta is None: + self._root_link_state_w_ta = ProxyArray(self._root_link_state_w.data) + return self._root_link_state_w_ta @property - def root_com_state_w(self) -> wp.array: + def root_com_state_w(self) -> ProxyArray: """Deprecated, same as :attr:`root_com_pose_w` and :attr:`root_com_vel_w`.""" warnings.warn( "The `root_com_state_w` property will be deprecated in IsaacLab 4.0. Please use `root_com_pose_w` and " @@ -842,10 +1033,12 @@ def root_com_state_w(self) -> wp.array: ) self._root_com_state_w.timestamp = self._sim_timestamp - return self._root_com_state_w.data + if self._root_com_state_w_ta is None: + self._root_com_state_w_ta = ProxyArray(self._root_com_state_w.data) + return self._root_com_state_w_ta @property - def body_state_w(self) -> wp.array: + def body_state_w(self) -> ProxyArray: """Deprecated, same as :attr:`body_link_pose_w` and :attr:`body_com_vel_w`.""" warnings.warn( "The `body_state_w` property will be deprecated in IsaacLab 4.0. Please use `body_link_pose_w` and " @@ -868,10 +1061,12 @@ def body_state_w(self) -> wp.array: device=self.device, ) self._root_state_w.timestamp = self._sim_timestamp - return self._root_state_w.data.reshape((self._num_instances, 1)) + if self._body_state_w_ta is None: + self._body_state_w_ta = ProxyArray(self._root_state_w.data.reshape((self._num_instances, 1))) + return self._body_state_w_ta @property - def body_link_state_w(self) -> wp.array: + def body_link_state_w(self) -> ProxyArray: """Deprecated, same as :attr:`body_link_pose_w` and :attr:`body_link_vel_w`.""" warnings.warn( "The `body_link_state_w` property will be deprecated in IsaacLab 4.0. Please use `body_link_pose_w` and " @@ -894,10 +1089,12 @@ def body_link_state_w(self) -> wp.array: device=self.device, ) self._root_link_state_w.timestamp = self._sim_timestamp - return self._root_link_state_w.data.reshape((self._num_instances, 1)) + if self._body_link_state_w_ta is None: + self._body_link_state_w_ta = ProxyArray(self._root_link_state_w.data.reshape((self._num_instances, 1))) + return self._body_link_state_w_ta @property - def body_com_state_w(self) -> wp.array: + def body_com_state_w(self) -> ProxyArray: """Deprecated, same as :attr:`body_com_pose_w` and :attr:`body_com_vel_w`.""" warnings.warn( "The `body_com_state_w` property will be deprecated in IsaacLab 4.0. Please use `body_com_pose_w` and " @@ -920,4 +1117,6 @@ def body_com_state_w(self) -> wp.array: device=self.device, ) self._root_com_state_w.timestamp = self._sim_timestamp - return self._root_com_state_w.data.reshape((self._num_instances, 1)) + if self._body_com_state_w_ta is None: + self._body_com_state_w_ta = ProxyArray(self._root_com_state_w.data.reshape((self._num_instances, 1))) + return self._body_com_state_w_ta 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 9d4d6b26979f..596af5c2bf11 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 @@ -192,8 +192,10 @@ def write_data_to_sim(self) -> None: composer = self._permanent_wrench_composer composer.compose_to_body_frame() self.root_view.apply_forces_and_torques_at_position( - force_data=self.reshape_data_to_view_2d(composer.out_force_b, device=self.device).view(wp.float32), - torque_data=self.reshape_data_to_view_2d(composer.out_torque_b, device=self.device).view(wp.float32), + force_data=self.reshape_data_to_view_2d(composer.out_force_b.warp, device=self.device).view(wp.float32), + torque_data=self.reshape_data_to_view_2d(composer.out_torque_b.warp, device=self.device).view( + wp.float32 + ), position_data=None, indices=self._env_body_ids_to_view_ids( self._ALL_ENV_INDICES, self._ALL_BODY_INDICES, device=self.device 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 dbb3b8523897..cee0d4de01dd 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 @@ -15,6 +15,7 @@ from isaaclab.assets.rigid_object_collection.base_rigid_object_collection_data import BaseRigidObjectCollectionData from isaaclab.utils.buffers import TimestampedBufferWarp as TimestampedBuffer from isaaclab.utils.math import normalize +from isaaclab.utils.warp import ProxyArray from isaaclab_physx.assets import kernels as shared_kernels from isaaclab_physx.physics import PhysxManager as SimulationManager @@ -45,6 +46,16 @@ class RigidObjectCollectionData(BaseRigidObjectCollectionData): The data is lazily updated, meaning that the data is only updated when it is accessed. This is useful when the data is expensive to compute or retrieve. The data is updated when the timestamp of the buffer is older than the current simulation timestamp. The timestamp is updated whenever the data is updated. + + .. note:: + **Pull-to-refresh model.** Properties pull fresh data from the PhysX tensor API on first + access per timestamp and cache the result. This differs from Newton, where buffers are + refreshed automatically by the simulation. + + .. note:: + **ProxyArray pointer stability.** Each :class:`ProxyArray` wrapper is created once and + reused because the PhysX tensor API returns views into stable, pre-allocated GPU buffers + whose device pointer does not change across simulation steps. """ __backend_name__: str = "physx" @@ -78,10 +89,14 @@ def __init__(self, root_view: physx.RigidBodyView, num_bodies: int, device: str) gravity_dir = normalize(gravity_dir.unsqueeze(0)).squeeze(0) # Initialize constants - self.GRAVITY_VEC_W = wp.from_torch(gravity_dir.repeat(self.num_instances, self.num_bodies, 1), dtype=wp.vec3f) - self.FORWARD_VEC_B = wp.from_torch( - torch.tensor((1.0, 0.0, 0.0), device=self.device).repeat(self.num_instances, self.num_bodies, 1), - dtype=wp.vec3f, + self.GRAVITY_VEC_W = ProxyArray( + wp.from_torch(gravity_dir.repeat(self.num_instances, self.num_bodies, 1), dtype=wp.vec3f) + ) + self.FORWARD_VEC_B = ProxyArray( + wp.from_torch( + torch.tensor((1.0, 0.0, 0.0), device=self.device).repeat(self.num_instances, self.num_bodies, 1), + dtype=wp.vec3f, + ) ) self._create_buffers() @@ -129,14 +144,16 @@ def update(self, dt: float) -> None: """ @property - def default_body_pose(self) -> wp.array: + def default_body_pose(self) -> ProxyArray: """Default body pose ``[pos, quat]`` in local environment frame. The position and quaternion are of the rigid body's actor frame. Shape is (num_instances, num_bodies), dtype = wp.transformf. In torch this resolves to (num_instances, num_bodies, 7). """ - return self._default_body_pose + if self._default_body_pose_ta is None: + self._default_body_pose_ta = ProxyArray(self._default_body_pose) + return self._default_body_pose_ta @default_body_pose.setter def default_body_pose(self, value: wp.array) -> None: @@ -153,14 +170,16 @@ def default_body_pose(self, value: wp.array) -> None: self._default_body_pose.assign(value) @property - def default_body_vel(self) -> wp.array: + def default_body_vel(self) -> ProxyArray: """Default body velocity ``[lin_vel, ang_vel]`` in local environment frame. The linear and angular velocities are of the rigid body's center of mass frame. Shape is (num_instances, num_bodies), dtype = wp.spatial_vectorf. In torch this resolves to (num_instances, num_bodies, 6). """ - return self._default_body_vel + if self._default_body_vel_ta is None: + self._default_body_vel_ta = ProxyArray(self._default_body_vel) + return self._default_body_vel_ta @default_body_vel.setter def default_body_vel(self, value: wp.array) -> None: @@ -181,7 +200,7 @@ def default_body_vel(self, value: wp.array) -> None: """ @property - def body_link_pose_w(self) -> wp.array: + def body_link_pose_w(self) -> ProxyArray: """Body link pose ``[pos, quat]`` in simulation world frame. Shape is (num_instances, num_bodies), dtype = wp.transformf. In torch this resolves to @@ -195,11 +214,18 @@ def body_link_pose_w(self) -> wp.array: # set the buffer data and timestamp self._body_link_pose_w.data = pose self._body_link_pose_w.timestamp = self._sim_timestamp + # Rebind ProxyArray since reshape creates a new wp.array each time + if self._body_link_pose_w_ta is not None: + self._body_link_pose_w_ta = ProxyArray(pose) + self._body_link_pos_w_ta = None + self._body_link_quat_w_ta = None - return self._body_link_pose_w.data + if self._body_link_pose_w_ta is None: + self._body_link_pose_w_ta = ProxyArray(self._body_link_pose_w.data) + return self._body_link_pose_w_ta @property - def body_link_vel_w(self) -> wp.array: + def body_link_vel_w(self) -> ProxyArray: """Body link velocity ``[lin_vel, ang_vel]`` in simulation world frame. Shape is (num_instances, num_bodies), dtype = wp.spatial_vectorf. In torch this resolves to @@ -223,10 +249,12 @@ def body_link_vel_w(self) -> wp.array: ) self._body_link_vel_w.timestamp = self._sim_timestamp - return self._body_link_vel_w.data + if self._body_link_vel_w_ta is None: + self._body_link_vel_w_ta = ProxyArray(self._body_link_vel_w.data) + return self._body_link_vel_w_ta @property - def body_com_pose_w(self) -> wp.array: + def body_com_pose_w(self) -> ProxyArray: """Body center of mass pose ``[pos, quat]`` in simulation world frame. Shape is (num_instances, num_bodies), dtype = wp.transformf. In torch this resolves to @@ -249,10 +277,12 @@ def body_com_pose_w(self) -> wp.array: ) self._body_com_pose_w.timestamp = self._sim_timestamp - return self._body_com_pose_w.data + if self._body_com_pose_w_ta is None: + self._body_com_pose_w_ta = ProxyArray(self._body_com_pose_w.data) + return self._body_com_pose_w_ta @property - def body_com_vel_w(self) -> wp.array: + def body_com_vel_w(self) -> ProxyArray: """Body center of mass velocity ``[lin_vel, ang_vel]`` in simulation world frame. Shape is (num_instances, num_bodies), dtype = wp.spatial_vectorf. In torch this resolves to @@ -264,11 +294,18 @@ def body_com_vel_w(self) -> wp.array: vel = self._reshape_view_to_data_2d(self._root_view.get_velocities().view(wp.spatial_vectorf)) self._body_com_vel_w.data = vel self._body_com_vel_w.timestamp = self._sim_timestamp + # Rebind ProxyArray since reshape creates a new wp.array each time + if self._body_com_vel_w_ta is not None: + self._body_com_vel_w_ta = ProxyArray(vel) + self._body_com_lin_vel_w_ta = None + self._body_com_ang_vel_w_ta = None - return self._body_com_vel_w.data + if self._body_com_vel_w_ta is None: + self._body_com_vel_w_ta = ProxyArray(self._body_com_vel_w.data) + return self._body_com_vel_w_ta @property - def body_com_acc_w(self) -> wp.array: + def body_com_acc_w(self) -> ProxyArray: """Acceleration of all bodies ``[lin_acc, ang_acc]`` in the simulation world frame. Shape is (num_instances, num_bodies), dtype = wp.spatial_vectorf. In torch this resolves to @@ -279,10 +316,17 @@ def body_com_acc_w(self) -> wp.array: acc = self._reshape_view_to_data_2d(self._root_view.get_accelerations().view(wp.spatial_vectorf)) self._body_com_acc_w.data = acc self._body_com_acc_w.timestamp = self._sim_timestamp - return self._body_com_acc_w.data + # Rebind ProxyArray since reshape creates a new wp.array each time + if self._body_com_acc_w_ta is not None: + self._body_com_acc_w_ta = ProxyArray(acc) + self._body_com_lin_acc_w_ta = None + self._body_com_ang_acc_w_ta = None + if self._body_com_acc_w_ta is None: + self._body_com_acc_w_ta = ProxyArray(self._body_com_acc_w.data) + return self._body_com_acc_w_ta @property - def body_com_pose_b(self) -> wp.array: + def body_com_pose_b(self) -> ProxyArray: """Center of mass pose ``[pos, quat]`` of all bodies in their respective body's link frames. Shape is (num_instances, num_bodies), dtype = wp.transformf. In torch this resolves to @@ -297,32 +341,38 @@ def body_com_pose_b(self) -> wp.array: self._body_com_pose_b.data.assign(poses) self._body_com_pose_b.timestamp = self._sim_timestamp - return self._body_com_pose_b.data + if self._body_com_pose_b_ta is None: + self._body_com_pose_b_ta = ProxyArray(self._body_com_pose_b.data) + return self._body_com_pose_b_ta @property - def body_mass(self) -> wp.array: + def body_mass(self) -> ProxyArray: """Mass of all bodies in the simulation world frame. Shape is (num_instances, num_bodies), dtype = wp.float32. In torch this resolves to (num_instances, num_bodies). """ - return self._body_mass + if self._body_mass_ta is None: + self._body_mass_ta = ProxyArray(self._body_mass) + return self._body_mass_ta @property - def body_inertia(self) -> wp.array: + def body_inertia(self) -> ProxyArray: """Inertia of all bodies in the simulation world frame. Shape is (num_instances, num_bodies, 9), dtype = wp.float32. In torch this resolves to (num_instances, num_bodies, 9). """ - return self._body_inertia + if self._body_inertia_ta is None: + self._body_inertia_ta = ProxyArray(self._body_inertia) + return self._body_inertia_ta """ Derived Properties. """ @property - def projected_gravity_b(self) -> wp.array: + def projected_gravity_b(self) -> ProxyArray: """Projection of the gravity direction on base frame. Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to @@ -337,10 +387,12 @@ def projected_gravity_b(self) -> wp.array: device=self.device, ) self._projected_gravity_b.timestamp = self._sim_timestamp - return self._projected_gravity_b.data + if self._projected_gravity_b_ta is None: + self._projected_gravity_b_ta = ProxyArray(self._projected_gravity_b.data) + return self._projected_gravity_b_ta @property - def heading_w(self) -> wp.array: + def heading_w(self) -> ProxyArray: """Yaw heading of the base frame (in radians). Shape is (num_instances, num_bodies), dtype = wp.float32. In torch this resolves to (num_instances, num_bodies). @@ -358,10 +410,12 @@ def heading_w(self) -> wp.array: device=self.device, ) self._heading_w.timestamp = self._sim_timestamp - return self._heading_w.data + if self._heading_w_ta is None: + self._heading_w_ta = ProxyArray(self._heading_w.data) + return self._heading_w_ta @property - def body_link_lin_vel_b(self) -> wp.array: + def body_link_lin_vel_b(self) -> ProxyArray: """Root link linear velocity in base frame. Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to @@ -378,10 +432,12 @@ def body_link_lin_vel_b(self) -> wp.array: device=self.device, ) self._body_link_lin_vel_b.timestamp = self._sim_timestamp - return self._body_link_lin_vel_b.data + if self._body_link_lin_vel_b_ta is None: + self._body_link_lin_vel_b_ta = ProxyArray(self._body_link_lin_vel_b.data) + return self._body_link_lin_vel_b_ta @property - def body_link_ang_vel_b(self) -> wp.array: + def body_link_ang_vel_b(self) -> ProxyArray: """Root link angular velocity in base frame. Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to @@ -398,10 +454,12 @@ def body_link_ang_vel_b(self) -> wp.array: device=self.device, ) self._body_link_ang_vel_b.timestamp = self._sim_timestamp - return self._body_link_ang_vel_b.data + if self._body_link_ang_vel_b_ta is None: + self._body_link_ang_vel_b_ta = ProxyArray(self._body_link_ang_vel_b.data) + return self._body_link_ang_vel_b_ta @property - def body_com_lin_vel_b(self) -> wp.array: + def body_com_lin_vel_b(self) -> ProxyArray: """Root center of mass linear velocity in base frame. Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to @@ -418,10 +476,12 @@ def body_com_lin_vel_b(self) -> wp.array: device=self.device, ) self._body_com_lin_vel_b.timestamp = self._sim_timestamp - return self._body_com_lin_vel_b.data + if self._body_com_lin_vel_b_ta is None: + self._body_com_lin_vel_b_ta = ProxyArray(self._body_com_lin_vel_b.data) + return self._body_com_lin_vel_b_ta @property - def body_com_ang_vel_b(self) -> wp.array: + def body_com_ang_vel_b(self) -> ProxyArray: """Root center of mass angular velocity in base frame. Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to @@ -438,124 +498,159 @@ def body_com_ang_vel_b(self) -> wp.array: device=self.device, ) self._body_com_ang_vel_b.timestamp = self._sim_timestamp - return self._body_com_ang_vel_b.data + if self._body_com_ang_vel_b_ta is None: + self._body_com_ang_vel_b_ta = ProxyArray(self._body_com_ang_vel_b.data) + return self._body_com_ang_vel_b_ta """ Sliced properties. """ @property - def body_link_pos_w(self) -> wp.array: + def body_link_pos_w(self) -> ProxyArray: """Positions of all bodies in simulation world frame. Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to (num_instances, num_bodies, 3). This quantity is the position of the rigid bodies' actor frame relative to the world. """ - return self._get_pos_from_transform(self.body_link_pose_w) + parent = self.body_link_pose_w + if self._body_link_pos_w_ta is None: + self._body_link_pos_w_ta = ProxyArray(self._get_pos_from_transform(parent.warp)) + return self._body_link_pos_w_ta @property - def body_link_quat_w(self) -> wp.array: + def body_link_quat_w(self) -> ProxyArray: """Orientation (x, y, z, w) of all bodies in simulation world frame. Shape is (num_instances, num_bodies), dtype = wp.quatf. In torch this resolves to (num_instances, num_bodies, 4). This quantity is the orientation of the rigid bodies' actor frame relative to the world. """ - return self._get_quat_from_transform(self.body_link_pose_w) + parent = self.body_link_pose_w + if self._body_link_quat_w_ta is None: + self._body_link_quat_w_ta = ProxyArray(self._get_quat_from_transform(parent.warp)) + return self._body_link_quat_w_ta @property - def body_link_lin_vel_w(self) -> wp.array: + def body_link_lin_vel_w(self) -> ProxyArray: """Linear velocity of all bodies in simulation world frame. Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to (num_instances, num_bodies, 3). This quantity is the linear velocity of the rigid bodies' actor frame relative to the world. """ - return self._get_lin_vel_from_spatial_vector(self.body_link_vel_w) + parent = self.body_link_vel_w + if self._body_link_lin_vel_w_ta is None: + self._body_link_lin_vel_w_ta = ProxyArray(self._get_lin_vel_from_spatial_vector(parent.warp)) + return self._body_link_lin_vel_w_ta @property - def body_link_ang_vel_w(self) -> wp.array: + def body_link_ang_vel_w(self) -> ProxyArray: """Angular velocity of all bodies in simulation world frame. Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to (num_instances, num_bodies, 3). This quantity is the angular velocity of the rigid bodies' actor frame relative to the world. """ - return self._get_ang_vel_from_spatial_vector(self.body_link_vel_w) + parent = self.body_link_vel_w + if self._body_link_ang_vel_w_ta is None: + self._body_link_ang_vel_w_ta = ProxyArray(self._get_ang_vel_from_spatial_vector(parent.warp)) + return self._body_link_ang_vel_w_ta @property - def body_com_pos_w(self) -> wp.array: + def body_com_pos_w(self) -> ProxyArray: """Positions of all bodies' center of mass in simulation world frame. Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to (num_instances, num_bodies, 3). This quantity is the position of the rigid bodies' center of mass frame. """ - return self._get_pos_from_transform(self.body_com_pose_w) + parent = self.body_com_pose_w + if self._body_com_pos_w_ta is None: + self._body_com_pos_w_ta = ProxyArray(self._get_pos_from_transform(parent.warp)) + return self._body_com_pos_w_ta @property - def body_com_quat_w(self) -> wp.array: + def body_com_quat_w(self) -> ProxyArray: """Orientation (x, y, z, w) of the principal axes of inertia of all bodies in simulation world frame. Shape is (num_instances, num_bodies), dtype = wp.quatf. In torch this resolves to (num_instances, num_bodies, 4). This quantity is the orientation of the principal axes of inertia of the rigid bodies. """ - return self._get_quat_from_transform(self.body_com_pose_w) + parent = self.body_com_pose_w + if self._body_com_quat_w_ta is None: + self._body_com_quat_w_ta = ProxyArray(self._get_quat_from_transform(parent.warp)) + return self._body_com_quat_w_ta @property - def body_com_lin_vel_w(self) -> wp.array: + def body_com_lin_vel_w(self) -> ProxyArray: """Linear velocity of all bodies in simulation world frame. Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to (num_instances, num_bodies, 3). This quantity is the linear velocity of the rigid bodies' center of mass frame. """ - return self._get_lin_vel_from_spatial_vector(self.body_com_vel_w) + parent = self.body_com_vel_w + if self._body_com_lin_vel_w_ta is None: + self._body_com_lin_vel_w_ta = ProxyArray(self._get_lin_vel_from_spatial_vector(parent.warp)) + return self._body_com_lin_vel_w_ta @property - def body_com_ang_vel_w(self) -> wp.array: + def body_com_ang_vel_w(self) -> ProxyArray: """Angular velocity of all bodies in simulation world frame. Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to (num_instances, num_bodies, 3). This quantity is the angular velocity of the rigid bodies' center of mass frame. """ - return self._get_ang_vel_from_spatial_vector(self.body_com_vel_w) + parent = self.body_com_vel_w + if self._body_com_ang_vel_w_ta is None: + self._body_com_ang_vel_w_ta = ProxyArray(self._get_ang_vel_from_spatial_vector(parent.warp)) + return self._body_com_ang_vel_w_ta @property - def body_com_lin_acc_w(self) -> wp.array: + def body_com_lin_acc_w(self) -> ProxyArray: """Linear acceleration of all bodies in simulation world frame. Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to (num_instances, num_bodies, 3). This quantity is the linear acceleration of the rigid bodies' center of mass frame. """ - return self._get_lin_vel_from_spatial_vector(self.body_com_acc_w) + parent = self.body_com_acc_w + if self._body_com_lin_acc_w_ta is None: + self._body_com_lin_acc_w_ta = ProxyArray(self._get_lin_vel_from_spatial_vector(parent.warp)) + return self._body_com_lin_acc_w_ta @property - def body_com_ang_acc_w(self) -> wp.array: + def body_com_ang_acc_w(self) -> ProxyArray: """Angular acceleration of all bodies in simulation world frame. Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to (num_instances, num_bodies, 3). This quantity is the angular acceleration of the rigid bodies' center of mass frame. """ - return self._get_ang_vel_from_spatial_vector(self.body_com_acc_w) + parent = self.body_com_acc_w + if self._body_com_ang_acc_w_ta is None: + self._body_com_ang_acc_w_ta = ProxyArray(self._get_ang_vel_from_spatial_vector(parent.warp)) + return self._body_com_ang_acc_w_ta @property - def body_com_pos_b(self) -> wp.array: + def body_com_pos_b(self) -> ProxyArray: """Center of mass position of all of the bodies in their respective link frames. Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to (num_instances, num_bodies, 3). This quantity is the center of mass location relative to its body's link frame. """ - return self._get_pos_from_transform(self.body_com_pose_b) + parent = self.body_com_pose_b + if self._body_com_pos_b_ta is None: + self._body_com_pos_b_ta = ProxyArray(self._get_pos_from_transform(parent.warp)) + return self._body_com_pos_b_ta @property - def body_com_quat_b(self) -> wp.array: + def body_com_quat_b(self) -> ProxyArray: """Orientation (x, y, z, w) of the principal axes of inertia of all of the bodies in their respective link frames. @@ -563,7 +658,10 @@ def body_com_quat_b(self) -> wp.array: (num_instances, num_bodies, 4). This quantity is the orientation of the principal axes of inertia relative to its body's link frame. """ - return self._get_quat_from_transform(self.body_com_pose_b) + parent = self.body_com_pose_b + if self._body_com_quat_b_ta is None: + self._body_com_quat_b_ta = ProxyArray(self._get_quat_from_transform(parent.warp)) + return self._body_com_quat_b_ta def _create_buffers(self) -> None: super()._create_buffers() @@ -616,6 +714,58 @@ def _create_buffers(self) -> None: # _reshape_view_to_data only handles single-element dtypes, so we use _reshape_view_to_data_3d. self._body_inertia = self._reshape_view_to_data_3d(self._root_view.get_inertias(), 9) + # Initialize ProxyArray wrappers + self._pin_proxy_arrays() + + def _pin_proxy_arrays(self) -> None: + """Create pinned ProxyArray wrappers for all data buffers. + + This is called once from :meth:`_create_buffers` during initialization. + PhysX tensor API buffers have stable GPU pointers across simulation steps, + so no rebinding is needed (unlike Newton). + """ + # -- Pinned ProxyArray cache (one per read property, lazily created on first access) + # Defaults + self._default_body_pose_ta: ProxyArray | None = None + self._default_body_vel_ta: ProxyArray | None = None + # Body state (timestamped) + self._body_link_pose_w_ta: ProxyArray | None = None + self._body_link_vel_w_ta: ProxyArray | None = None + self._body_com_pose_w_ta: ProxyArray | None = None + self._body_com_vel_w_ta: ProxyArray | None = None + self._body_com_acc_w_ta: ProxyArray | None = None + self._body_com_pose_b_ta: ProxyArray | None = None + # Body properties + self._body_mass_ta: ProxyArray | None = None + self._body_inertia_ta: ProxyArray | None = None + # Derived properties (timestamped) + self._projected_gravity_b_ta: ProxyArray | None = None + self._heading_w_ta: ProxyArray | None = None + self._body_link_lin_vel_b_ta: ProxyArray | None = None + self._body_link_ang_vel_b_ta: ProxyArray | None = None + self._body_com_lin_vel_b_ta: ProxyArray | None = None + self._body_com_ang_vel_b_ta: ProxyArray | None = None + # Sliced properties (body link) + self._body_link_pos_w_ta: ProxyArray | None = None + self._body_link_quat_w_ta: ProxyArray | None = None + self._body_link_lin_vel_w_ta: ProxyArray | None = None + self._body_link_ang_vel_w_ta: ProxyArray | None = None + # Sliced properties (body com) + self._body_com_pos_w_ta: ProxyArray | None = None + self._body_com_quat_w_ta: ProxyArray | None = None + self._body_com_lin_vel_w_ta: ProxyArray | None = None + self._body_com_ang_vel_w_ta: ProxyArray | None = None + self._body_com_lin_acc_w_ta: ProxyArray | None = None + self._body_com_ang_acc_w_ta: ProxyArray | None = None + # Sliced properties (body com in body frame) + self._body_com_pos_b_ta: ProxyArray | None = None + self._body_com_quat_b_ta: ProxyArray | None = None + # Deprecated state-concat properties + self._default_body_state_ta: ProxyArray | None = None + self._body_state_w_ta: ProxyArray | None = None + self._body_link_state_w_ta: ProxyArray | None = None + self._body_com_state_w_ta: ProxyArray | None = None + """ Helpers. """ @@ -669,7 +819,7 @@ def _reshape_view_to_data_3d(self, data: wp.array, data_dim: int) -> wp.array: ) return wp.clone(strided_view, self.device) - def _get_pos_from_transform(self, transform: wp.array) -> wp.array: + def _get_pos_from_transform(self, transform: wp.array) -> ProxyArray: """Generates a position array from a transform array.""" return wp.array( ptr=transform.ptr, @@ -679,7 +829,7 @@ def _get_pos_from_transform(self, transform: wp.array) -> wp.array: device=self.device, ) - def _get_quat_from_transform(self, transform: wp.array) -> wp.array: + def _get_quat_from_transform(self, transform: wp.array) -> ProxyArray: """Generates a quaternion array from a transform array.""" return wp.array( ptr=transform.ptr + 3 * 4, @@ -689,7 +839,7 @@ def _get_quat_from_transform(self, transform: wp.array) -> wp.array: device=self.device, ) - def _get_lin_vel_from_spatial_vector(self, spatial_vector: wp.array) -> wp.array: + def _get_lin_vel_from_spatial_vector(self, spatial_vector: wp.array) -> ProxyArray: """Generates a linear velocity array from a spatial vector array.""" return wp.array( ptr=spatial_vector.ptr, @@ -699,7 +849,7 @@ def _get_lin_vel_from_spatial_vector(self, spatial_vector: wp.array) -> wp.array device=self.device, ) - def _get_ang_vel_from_spatial_vector(self, spatial_vector: wp.array) -> wp.array: + def _get_ang_vel_from_spatial_vector(self, spatial_vector: wp.array) -> ProxyArray: """Generates an angular velocity array from a spatial vector array.""" return wp.array( ptr=spatial_vector.ptr + 3 * 4, @@ -714,7 +864,7 @@ def _get_ang_vel_from_spatial_vector(self, spatial_vector: wp.array) -> wp.array """ @property - def default_body_state(self) -> wp.array: + def default_body_state(self) -> ProxyArray: """Default root state ``[pos, quat, lin_vel, ang_vel]`` in local environment frame. The position and quaternion are of the rigid body's actor frame. Meanwhile, the linear and angular velocities @@ -742,10 +892,12 @@ def default_body_state(self) -> wp.array: ], device=self.device, ) - return self._default_body_state + if self._default_body_state_ta is None: + self._default_body_state_ta = ProxyArray(self._default_body_state) + return self._default_body_state_ta @property - def body_state_w(self) -> wp.array: + def body_state_w(self) -> ProxyArray: """Deprecated, same as :attr:`body_link_pose_w` and :attr:`body_com_vel_w`.""" warnings.warn( "The `body_state_w` property will be deprecated in IsaacLab 4.0. Please use `body_link_pose_w` and " @@ -768,10 +920,12 @@ def body_state_w(self) -> wp.array: ) self._body_state_w.timestamp = self._sim_timestamp - return self._body_state_w.data + if self._body_state_w_ta is None: + self._body_state_w_ta = ProxyArray(self._body_state_w.data) + return self._body_state_w_ta @property - def body_link_state_w(self) -> wp.array: + def body_link_state_w(self) -> ProxyArray: """Deprecated, same as :attr:`body_link_pose_w` and :attr:`body_link_vel_w`.""" warnings.warn( "The `body_link_state_w` property will be deprecated in IsaacLab 4.0. Please use `body_link_pose_w` and " @@ -794,10 +948,12 @@ def body_link_state_w(self) -> wp.array: ) self._body_link_state_w.timestamp = self._sim_timestamp - return self._body_link_state_w.data + if self._body_link_state_w_ta is None: + self._body_link_state_w_ta = ProxyArray(self._body_link_state_w.data) + return self._body_link_state_w_ta @property - def body_com_state_w(self) -> wp.array: + def body_com_state_w(self) -> ProxyArray: """Deprecated, same as :attr:`body_com_pose_w` and :attr:`body_com_vel_w`.""" warnings.warn( "The `body_com_state_w` property will be deprecated in IsaacLab 4.0. Please use `body_com_pose_w` and " @@ -820,4 +976,6 @@ def body_com_state_w(self) -> wp.array: ) self._body_com_state_w.timestamp = self._sim_timestamp - return self._body_com_state_w.data + if self._body_com_state_w_ta is None: + self._body_com_state_w_ta = ProxyArray(self._body_com_state_w.data) + return self._body_com_state_w_ta diff --git a/source/isaaclab_physx/isaaclab_physx/scene_data_providers/physx_scene_data_provider.py b/source/isaaclab_physx/isaaclab_physx/scene_data_providers/physx_scene_data_provider.py index fde9402c5f2e..aa19784ff1b6 100644 --- a/source/isaaclab_physx/isaaclab_physx/scene_data_providers/physx_scene_data_provider.py +++ b/source/isaaclab_physx/isaaclab_physx/scene_data_providers/physx_scene_data_provider.py @@ -486,10 +486,10 @@ def _apply_xform_poses(self, positions: Any, orientations: Any, covered: Any, xf path, device=self._device, stage=self._stage, validate_xform_ops=False ) - pos_wp, quat_wp = self._xform_views[path].get_world_poses() - if pos_wp is not None and quat_wp is not None: - positions[idx] = wp.to_torch(pos_wp).to(device=self._device, dtype=torch.float32).squeeze() - orientations[idx] = wp.to_torch(quat_wp).to(device=self._device, dtype=torch.float32).squeeze() + pos_w, quat_w = self._xform_views[path].get_world_poses() + if pos_w is not None and quat_w is not None: + positions[idx] = pos_w.torch.to(device=self._device, dtype=torch.float32).squeeze() + orientations[idx] = quat_w.torch.to(device=self._device, dtype=torch.float32).squeeze() covered[idx] = True xform_mask[idx] = True count += 1 diff --git a/source/isaaclab_physx/isaaclab_physx/sensors/contact_sensor/contact_sensor.py b/source/isaaclab_physx/isaaclab_physx/sensors/contact_sensor/contact_sensor.py index 16f7046dbe97..c38a0fffe281 100644 --- a/source/isaaclab_physx/isaaclab_physx/sensors/contact_sensor/contact_sensor.py +++ b/source/isaaclab_physx/isaaclab_physx/sensors/contact_sensor/contact_sensor.py @@ -20,6 +20,7 @@ from isaaclab.app.settings_manager import get_settings_manager from isaaclab.markers import VisualizationMarkers from isaaclab.sensors.contact_sensor import BaseContactSensor +from isaaclab.utils.warp import ProxyArray from isaaclab_physx.physics import PhysxManager as SimulationManager @@ -189,7 +190,7 @@ def reset(self, env_ids: Sequence[int] | None = None, env_mask: wp.array | None device=self._device, ) - def compute_first_contact(self, dt: float, abs_tol: float = 1.0e-8) -> wp.array: + def compute_first_contact(self, dt: float, abs_tol: float = 1.0e-8) -> ProxyArray: """Checks if bodies that have established contact within the last :attr:`dt` seconds. This function checks if the bodies have established contact within the last :attr:`dt` seconds @@ -232,9 +233,9 @@ def compute_first_contact(self, dt: float, abs_tol: float = 1.0e-8) -> wp.array: outputs=[self._data._first_transition], device=self._device, ) - return self._data._first_transition + return self._data._first_transition_ta - def compute_first_air(self, dt: float, abs_tol: float = 1.0e-8) -> wp.array: + def compute_first_air(self, dt: float, abs_tol: float = 1.0e-8) -> ProxyArray: """Checks if bodies that have broken contact within the last :attr:`dt` seconds. This function checks if the bodies have broken contact within the last :attr:`dt` seconds @@ -277,7 +278,7 @@ def compute_first_air(self, dt: float, abs_tol: float = 1.0e-8) -> wp.array: outputs=[self._data._first_transition], device=self._device, ) - return self._data._first_transition + return self._data._first_transition_ta """ Implementation. @@ -478,13 +479,13 @@ def _debug_vis_callback(self, event): if self.body_physx_view is None: return # Convert warp data to torch at the boundary for visualization - net_forces_torch = wp.to_torch(self._data._net_forces_w) # (N, B, 3) + net_forces_torch = self._data.net_forces_w.torch # (N, B, 3) net_contact_force_w = torch.linalg.norm(net_forces_torch, dim=-1) # marker indices: 0 = contact, 1 = no contact marker_indices = torch.where(net_contact_force_w > self.cfg.force_threshold, 0, 1) # check if prim is visualized if self.cfg.track_pose: - frame_origins = wp.to_torch(self._data._pos_w) # (N, B, 3) + frame_origins = self._data.pos_w.torch # (N, B, 3) else: pose = self.body_physx_view.get_transforms() # (N*B, 7) float32 pose_torch = wp.to_torch(pose) diff --git a/source/isaaclab_physx/isaaclab_physx/sensors/contact_sensor/contact_sensor_data.py b/source/isaaclab_physx/isaaclab_physx/sensors/contact_sensor/contact_sensor_data.py index 93f12c33918b..5e50c42da7ef 100644 --- a/source/isaaclab_physx/isaaclab_physx/sensors/contact_sensor/contact_sensor_data.py +++ b/source/isaaclab_physx/isaaclab_physx/sensors/contact_sensor/contact_sensor_data.py @@ -11,6 +11,7 @@ import warp as wp from isaaclab.sensors.contact_sensor import BaseContactSensorData +from isaaclab.utils.warp import ProxyArray from isaaclab_physx.sensors.kernels import concat_pos_and_quat_to_pose_kernel @@ -21,11 +22,13 @@ class ContactSensorData(BaseContactSensorData): """Data container for the PhysX contact reporting sensor.""" @property - def pose_w(self) -> wp.array | None: + def pose_w(self) -> ProxyArray | None: """Pose of the sensor origin in world frame. None if :attr:`ContactSensorCfg.track_pose` is False. """ + if self._pose_w is None: + return None wp.launch( concat_pos_and_quat_to_pose_kernel, dim=(self._num_envs, self._num_sensors), @@ -33,10 +36,12 @@ def pose_w(self) -> wp.array | None: outputs=[self._pose_w], device=self._device, ) - return self._pose_w + if self._pose_w_ta is None: + self._pose_w_ta = ProxyArray(self._pose_w) + return self._pose_w_ta @property - def pos_w(self) -> wp.array | None: + def pos_w(self) -> ProxyArray | None: """Position of the sensor origin in world frame. Shape is (num_instances, num_sensors), dtype = wp.vec3f. In torch this resolves to @@ -44,10 +49,14 @@ def pos_w(self) -> wp.array | None: None if :attr:`ContactSensorCfg.track_pose` is False. """ - return self._pos_w + if self._pos_w is None: + return None + if self._pos_w_ta is None: + self._pos_w_ta = ProxyArray(self._pos_w) + return self._pos_w_ta @property - def quat_w(self) -> wp.array | None: + def quat_w(self) -> ProxyArray | None: """Orientation of the sensor origin in world frame. Shape is (num_instances, num_sensors), dtype = wp.quatf. In torch this resolves to @@ -55,28 +64,40 @@ def quat_w(self) -> wp.array | None: None if :attr:`ContactSensorCfg.track_pose` is False. """ - return self._quat_w + if self._quat_w is None: + return None + if self._quat_w_ta is None: + self._quat_w_ta = ProxyArray(self._quat_w) + return self._quat_w_ta @property - def net_forces_w(self) -> wp.array | None: + def net_forces_w(self) -> ProxyArray | None: """The net normal contact forces in world frame. Shape is (num_instances, num_sensors), dtype = wp.vec3f. In torch this resolves to (num_instances, num_sensors, 3). """ - return self._net_forces_w + if self._net_forces_w is None: + return None + if self._net_forces_w_ta is None: + self._net_forces_w_ta = ProxyArray(self._net_forces_w) + return self._net_forces_w_ta @property - def net_forces_w_history(self) -> wp.array | None: + def net_forces_w_history(self) -> ProxyArray | None: """History of net normal contact forces. Shape is (num_instances, history_length, num_sensors), dtype = wp.vec3f. In torch this resolves to (num_instances, history_length, num_sensors, 3). """ - return self._net_forces_w_history + if self._net_forces_w_history is None: + return None + if self._net_forces_w_history_ta is None: + self._net_forces_w_history_ta = ProxyArray(self._net_forces_w_history) + return self._net_forces_w_history_ta @property - def force_matrix_w(self) -> wp.array | None: + def force_matrix_w(self) -> ProxyArray | None: """Normal contact forces filtered between sensor and filtered bodies. Shape is (num_instances, num_sensors, num_filter_shapes), dtype = wp.vec3f. In torch this resolves to @@ -84,10 +105,14 @@ def force_matrix_w(self) -> wp.array | None: None if :attr:`ContactSensorCfg.filter_prim_paths_expr` is empty. """ - return self._force_matrix_w + if self._force_matrix_w is None: + return None + if self._force_matrix_w_ta is None: + self._force_matrix_w_ta = ProxyArray(self._force_matrix_w) + return self._force_matrix_w_ta @property - def force_matrix_w_history(self) -> wp.array | None: + def force_matrix_w_history(self) -> ProxyArray | None: """History of filtered contact forces. Shape is (num_instances, history_length, num_sensors, num_filter_shapes), dtype = wp.vec3f. @@ -95,10 +120,14 @@ def force_matrix_w_history(self) -> wp.array | None: None if :attr:`ContactSensorCfg.filter_prim_paths_expr` is empty. """ - return self._force_matrix_w_history + if self._force_matrix_w_history is None: + return None + if self._force_matrix_w_history_ta is None: + self._force_matrix_w_history_ta = ProxyArray(self._force_matrix_w_history) + return self._force_matrix_w_history_ta @property - def contact_pos_w(self) -> wp.array | None: + def contact_pos_w(self) -> ProxyArray | None: """Average position of contact points. Shape is (num_instances, num_sensors, num_filter_shapes), dtype = wp.vec3f. In torch this resolves to @@ -106,10 +135,14 @@ def contact_pos_w(self) -> wp.array | None: None if :attr:`ContactSensorCfg.track_contact_points` is False. """ - return self._contact_pos_w + if self._contact_pos_w is None: + return None + if self._contact_pos_w_ta is None: + self._contact_pos_w_ta = ProxyArray(self._contact_pos_w) + return self._contact_pos_w_ta @property - def friction_forces_w(self) -> wp.array | None: + def friction_forces_w(self) -> ProxyArray | None: """Sum of friction forces. Shape is (num_instances, num_sensors, num_filter_shapes), dtype = wp.vec3f. In torch this resolves to @@ -117,47 +150,67 @@ def friction_forces_w(self) -> wp.array | None: None if :attr:`ContactSensorCfg.track_friction_forces` is False. """ - return self._friction_forces_w + if self._friction_forces_w is None: + return None + if self._friction_forces_w_ta is None: + self._friction_forces_w_ta = ProxyArray(self._friction_forces_w) + return self._friction_forces_w_ta @property - def last_air_time(self) -> wp.array | None: + def last_air_time(self) -> ProxyArray | None: """Time spent in air before last contact. Shape is (num_instances, num_sensors), dtype = wp.float32. None if :attr:`ContactSensorCfg.track_air_time` is False. """ - return self._last_air_time + if self._last_air_time is None: + return None + if self._last_air_time_ta is None: + self._last_air_time_ta = ProxyArray(self._last_air_time) + return self._last_air_time_ta @property - def current_air_time(self) -> wp.array | None: + def current_air_time(self) -> ProxyArray | None: """Time spent in air since last detach. Shape is (num_instances, num_sensors), dtype = wp.float32. None if :attr:`ContactSensorCfg.track_air_time` is False. """ - return self._current_air_time + if self._current_air_time is None: + return None + if self._current_air_time_ta is None: + self._current_air_time_ta = ProxyArray(self._current_air_time) + return self._current_air_time_ta @property - def last_contact_time(self) -> wp.array | None: + def last_contact_time(self) -> ProxyArray | None: """Time spent in contact before last detach. Shape is (num_instances, num_sensors), dtype = wp.float32. None if :attr:`ContactSensorCfg.track_air_time` is False. """ - return self._last_contact_time + if self._last_contact_time is None: + return None + if self._last_contact_time_ta is None: + self._last_contact_time_ta = ProxyArray(self._last_contact_time) + return self._last_contact_time_ta @property - def current_contact_time(self) -> wp.array | None: + def current_contact_time(self) -> ProxyArray | None: """Time spent in contact since last contact. Shape is (num_instances, num_sensors), dtype = wp.float32. None if :attr:`ContactSensorCfg.track_air_time` is False. """ - return self._current_contact_time + if self._current_contact_time is None: + return None + if self._current_contact_time_ta is None: + self._current_contact_time_ta = ProxyArray(self._current_contact_time) + return self._current_contact_time_ta def create_buffers( self, @@ -221,12 +274,14 @@ def create_buffers( self._last_contact_time = wp.zeros((num_envs, num_sensors), dtype=wp.float32, device=device) self._current_contact_time = wp.zeros((num_envs, num_sensors), dtype=wp.float32, device=device) self._first_transition = wp.zeros((num_envs, num_sensors), dtype=wp.float32, device=device) + self._first_transition_ta = ProxyArray(self._first_transition) else: self._last_air_time = None self._current_air_time = None self._last_contact_time = None self._current_contact_time = None self._first_transition = None + self._first_transition_ta = None # Track contact points if requested - filled with NaN if track_contact_points: @@ -246,3 +301,18 @@ def create_buffers( ) else: self._friction_forces_w = None + + # -- Pinned ProxyArray cache (one per read property, lazily created on first access) + self._pose_w_ta: ProxyArray | None = None + self._pos_w_ta: ProxyArray | None = None + self._quat_w_ta: ProxyArray | None = None + self._net_forces_w_ta: ProxyArray | None = None + self._net_forces_w_history_ta: ProxyArray | None = None + self._force_matrix_w_ta: ProxyArray | None = None + self._force_matrix_w_history_ta: ProxyArray | None = None + self._contact_pos_w_ta: ProxyArray | None = None + self._friction_forces_w_ta: ProxyArray | None = None + self._last_air_time_ta: ProxyArray | None = None + self._current_air_time_ta: ProxyArray | None = None + self._last_contact_time_ta: ProxyArray | None = None + self._current_contact_time_ta: ProxyArray | None = None diff --git a/source/isaaclab_physx/isaaclab_physx/sensors/frame_transformer/frame_transformer_data.py b/source/isaaclab_physx/isaaclab_physx/sensors/frame_transformer/frame_transformer_data.py index f73c00bbfaf2..7f7c0a837b2f 100644 --- a/source/isaaclab_physx/isaaclab_physx/sensors/frame_transformer/frame_transformer_data.py +++ b/source/isaaclab_physx/isaaclab_physx/sensors/frame_transformer/frame_transformer_data.py @@ -8,6 +8,7 @@ import warp as wp from isaaclab.sensors.frame_transformer import BaseFrameTransformerData +from isaaclab.utils.warp import ProxyArray from isaaclab_physx.sensors.kernels import concat_pos_and_quat_to_pose_1d_kernel, concat_pos_and_quat_to_pose_kernel @@ -21,7 +22,7 @@ def target_frame_names(self) -> list[str]: return self._target_frame_names @property - def target_pose_source(self) -> wp.array: + def target_pose_source(self) -> ProxyArray: """Pose of target frame(s) relative to source frame. Shape is (num_instances, num_target_frames), dtype = wp.transformf. In torch this resolves to @@ -34,28 +35,34 @@ def target_pose_source(self) -> wp.array: outputs=[self._target_pose_source], device=self._device, ) - return self._target_pose_source + if self._target_pose_source_ta is None: + self._target_pose_source_ta = ProxyArray(self._target_pose_source) + return self._target_pose_source_ta @property - def target_pos_source(self) -> wp.array: + def target_pos_source(self) -> ProxyArray: """Position of target frame(s) relative to source frame. Shape is (num_instances, num_target_frames), dtype = wp.vec3f. In torch this resolves to (num_instances, num_target_frames, 3). """ - return self._target_pos_source + if self._target_pos_source_ta is None: + self._target_pos_source_ta = ProxyArray(self._target_pos_source) + return self._target_pos_source_ta @property - def target_quat_source(self) -> wp.array: + def target_quat_source(self) -> ProxyArray: """Orientation of target frame(s) relative to source frame. Shape is (num_instances, num_target_frames), dtype = wp.quatf. In torch this resolves to (num_instances, num_target_frames, 4). The orientation is provided in (x, y, z, w) format. """ - return self._target_quat_source + if self._target_quat_source_ta is None: + self._target_quat_source_ta = ProxyArray(self._target_quat_source) + return self._target_quat_source_ta @property - def target_pose_w(self) -> wp.array: + def target_pose_w(self) -> ProxyArray: """Pose of target frame(s) after offset in world frame. Shape is (num_instances, num_target_frames), dtype = wp.transformf. In torch this resolves to @@ -68,28 +75,34 @@ def target_pose_w(self) -> wp.array: outputs=[self._target_pose_w], device=self._device, ) - return self._target_pose_w + if self._target_pose_w_ta is None: + self._target_pose_w_ta = ProxyArray(self._target_pose_w) + return self._target_pose_w_ta @property - def target_pos_w(self) -> wp.array: + def target_pos_w(self) -> ProxyArray: """Position of target frame(s) after offset in world frame. Shape is (num_instances, num_target_frames), dtype = wp.vec3f. In torch this resolves to (num_instances, num_target_frames, 3). """ - return self._target_pos_w + if self._target_pos_w_ta is None: + self._target_pos_w_ta = ProxyArray(self._target_pos_w) + return self._target_pos_w_ta @property - def target_quat_w(self) -> wp.array: + def target_quat_w(self) -> ProxyArray: """Orientation of target frame(s) after offset in world frame. Shape is (num_instances, num_target_frames), dtype = wp.quatf. In torch this resolves to (num_instances, num_target_frames, 4). The orientation is provided in (x, y, z, w) format. """ - return self._target_quat_w + if self._target_quat_w_ta is None: + self._target_quat_w_ta = ProxyArray(self._target_quat_w) + return self._target_quat_w_ta @property - def source_pose_w(self) -> wp.array: + def source_pose_w(self) -> ProxyArray: """Pose of source frame after offset in world frame. Shape is (num_instances,), dtype = wp.transformf. In torch this resolves to (num_instances, 7). @@ -102,24 +115,30 @@ def source_pose_w(self) -> wp.array: outputs=[self._source_pose_w], device=self._device, ) - return self._source_pose_w + if self._source_pose_w_ta is None: + self._source_pose_w_ta = ProxyArray(self._source_pose_w) + return self._source_pose_w_ta @property - def source_pos_w(self) -> wp.array: + def source_pos_w(self) -> ProxyArray: """Position of source frame after offset in world frame. Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). """ - return self._source_pos_w + if self._source_pos_w_ta is None: + self._source_pos_w_ta = ProxyArray(self._source_pos_w) + return self._source_pos_w_ta @property - def source_quat_w(self) -> wp.array: + def source_quat_w(self) -> ProxyArray: """Orientation of source frame after offset in world frame. Shape is (num_instances,), dtype = wp.quatf. In torch this resolves to (num_instances, 4). The orientation is provided in (x, y, z, w) format. """ - return self._source_quat_w + if self._source_quat_w_ta is None: + self._source_quat_w_ta = ProxyArray(self._source_quat_w) + return self._source_quat_w_ta def create_buffers( self, @@ -155,3 +174,14 @@ def create_buffers( wp.to_torch(self._source_quat_w)[:, 3] = 1.0 wp.to_torch(self._target_quat_w)[:, :, 3] = 1.0 wp.to_torch(self._target_quat_source)[:, :, 3] = 1.0 + + # -- Pinned ProxyArray cache (one per read property, lazily created on first access) + self._target_pose_source_ta: ProxyArray | None = None + self._target_pos_source_ta: ProxyArray | None = None + self._target_quat_source_ta: ProxyArray | None = None + self._target_pose_w_ta: ProxyArray | None = None + self._target_pos_w_ta: ProxyArray | None = None + self._target_quat_w_ta: ProxyArray | None = None + self._source_pose_w_ta: ProxyArray | None = None + self._source_pos_w_ta: ProxyArray | None = None + self._source_quat_w_ta: ProxyArray | None = None diff --git a/source/isaaclab_physx/isaaclab_physx/sensors/imu/imu_data.py b/source/isaaclab_physx/isaaclab_physx/sensors/imu/imu_data.py index ee76cd04ffe4..7251111c3eb5 100644 --- a/source/isaaclab_physx/isaaclab_physx/sensors/imu/imu_data.py +++ b/source/isaaclab_physx/isaaclab_physx/sensors/imu/imu_data.py @@ -8,28 +8,33 @@ import warp as wp from isaaclab.sensors.imu import BaseImuData +from isaaclab.utils.warp import ProxyArray class ImuData(BaseImuData): """Data container for the PhysX IMU sensor.""" @property - def ang_vel_b(self) -> wp.array: + def ang_vel_b(self) -> ProxyArray: """IMU frame angular velocity relative to the world expressed in IMU frame [rad/s]. Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). """ - return self._ang_vel_b + if self._ang_vel_b_ta is None: + self._ang_vel_b_ta = ProxyArray(self._ang_vel_b) + return self._ang_vel_b_ta @property - def lin_acc_b(self) -> wp.array: + def lin_acc_b(self) -> ProxyArray: """Linear acceleration (proper) in the IMU frame [m/s^2]. Zero in freefall, +g upward at rest. Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). """ - return self._lin_acc_b + if self._lin_acc_b_ta is None: + self._lin_acc_b_ta = ProxyArray(self._lin_acc_b) + return self._lin_acc_b_ta def create_buffers(self, num_envs: int, device: str) -> None: """Create internal buffers for sensor data. @@ -42,3 +47,7 @@ def create_buffers(self, num_envs: int, device: str) -> None: self._device = device self._ang_vel_b = wp.zeros(num_envs, dtype=wp.vec3f, device=device) self._lin_acc_b = wp.zeros(num_envs, dtype=wp.vec3f, device=device) + + # -- Pinned ProxyArray cache (one per read property, lazily created on first access) + self._ang_vel_b_ta: ProxyArray | None = None + self._lin_acc_b_ta: ProxyArray | None = None diff --git a/source/isaaclab_physx/isaaclab_physx/sensors/pva/pva.py b/source/isaaclab_physx/isaaclab_physx/sensors/pva/pva.py index 25e3fad0325d..32f4549e416d 100644 --- a/source/isaaclab_physx/isaaclab_physx/sensors/pva/pva.py +++ b/source/isaaclab_physx/isaaclab_physx/sensors/pva/pva.py @@ -17,6 +17,7 @@ import isaaclab.utils.math as math_utils from isaaclab.markers import VisualizationMarkers from isaaclab.sensors.pva import BasePva +from isaaclab.utils.warp import ProxyArray from isaaclab_physx.physics import PhysxManager as SimulationManager @@ -180,7 +181,7 @@ def _initialize_impl(self): 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_repeated = gravity_dir.repeat(self.num_instances, 1) - self.GRAVITY_VEC_W = wp.from_torch(gravity_dir_repeated.contiguous(), dtype=wp.vec3f) + self.GRAVITY_VEC_W = ProxyArray(wp.from_torch(gravity_dir_repeated.contiguous(), dtype=wp.vec3f)) # Create internal buffers self._initialize_buffers_impl() @@ -276,19 +277,17 @@ def _debug_vis_callback(self, event): return # get marker location # -- base state (convert warp -> torch for visualization) - base_pos_w = wp.to_torch(self._data.pos_w).clone() + base_pos_w = self._data.pos_w.torch.clone() base_pos_w[:, 2] += 0.5 # -- resolve the scales default_scale = self.acceleration_visualizer.cfg.markers["arrow"].scale - arrow_scale = torch.tensor(default_scale, device=self.device).repeat( - wp.to_torch(self._data.lin_acc_b).shape[0], 1 - ) + arrow_scale = torch.tensor(default_scale, device=self.device).repeat(self._data.lin_acc_b.torch.shape[0], 1) # get up axis of current stage up_axis = UsdGeom.GetStageUpAxis(self.stage) # arrow-direction - pos_w_torch = wp.to_torch(self._data.pos_w) - quat_w_torch = wp.to_torch(self._data.quat_w) - lin_acc_b_torch = wp.to_torch(self._data.lin_acc_b) + pos_w_torch = self._data.pos_w.torch + quat_w_torch = self._data.quat_w.torch + lin_acc_b_torch = self._data.lin_acc_b.torch quat_opengl = math_utils.quat_from_matrix( math_utils.create_rotation_matrix_from_view( pos_w_torch, diff --git a/source/isaaclab_physx/isaaclab_physx/sensors/pva/pva_data.py b/source/isaaclab_physx/isaaclab_physx/sensors/pva/pva_data.py index 07ace050afbd..3128551c3273 100644 --- a/source/isaaclab_physx/isaaclab_physx/sensors/pva/pva_data.py +++ b/source/isaaclab_physx/isaaclab_physx/sensors/pva/pva_data.py @@ -10,6 +10,7 @@ import warp as wp from isaaclab.sensors.pva import BasePvaData +from isaaclab.utils.warp import ProxyArray from isaaclab_physx.sensors.kernels import concat_pos_and_quat_to_pose_1d_kernel @@ -20,7 +21,7 @@ class PvaData(BasePvaData): """Data container for the PhysX PVA sensor.""" @property - def pose_w(self) -> wp.array: + def pose_w(self) -> ProxyArray: """Pose of the sensor origin in world frame. Shape is (num_instances,), dtype = wp.transformf. In torch this resolves to (num_instances, 7). @@ -33,66 +34,82 @@ def pose_w(self) -> wp.array: outputs=[self._pose_w], device=self._device, ) - return self._pose_w + if self._pose_w_ta is None: + self._pose_w_ta = ProxyArray(self._pose_w) + return self._pose_w_ta @property - def pos_w(self) -> wp.array: + def pos_w(self) -> ProxyArray: """Position of the sensor origin in world frame. Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). """ - return self._pos_w + if self._pos_w_ta is None: + self._pos_w_ta = ProxyArray(self._pos_w) + return self._pos_w_ta @property - def quat_w(self) -> wp.array: + def quat_w(self) -> ProxyArray: """Orientation of the sensor origin in world frame. Shape is (num_instances,), dtype = wp.quatf. In torch this resolves to (num_instances, 4). The orientation is provided in (x, y, z, w) format. """ - return self._quat_w + if self._quat_w_ta is None: + self._quat_w_ta = ProxyArray(self._quat_w) + return self._quat_w_ta @property - def projected_gravity_b(self) -> wp.array: + def projected_gravity_b(self) -> ProxyArray: """Gravity direction unit vector projected on the PVA frame. Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). """ - return self._projected_gravity_b + if self._projected_gravity_b_ta is None: + self._projected_gravity_b_ta = ProxyArray(self._projected_gravity_b) + return self._projected_gravity_b_ta @property - def lin_vel_b(self) -> wp.array: + def lin_vel_b(self) -> ProxyArray: """PVA frame linear velocity relative to the world expressed in PVA frame. Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). """ - return self._lin_vel_b + if self._lin_vel_b_ta is None: + self._lin_vel_b_ta = ProxyArray(self._lin_vel_b) + return self._lin_vel_b_ta @property - def ang_vel_b(self) -> wp.array: + def ang_vel_b(self) -> ProxyArray: """PVA frame angular velocity relative to the world expressed in PVA frame. Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). """ - return self._ang_vel_b + if self._ang_vel_b_ta is None: + self._ang_vel_b_ta = ProxyArray(self._ang_vel_b) + return self._ang_vel_b_ta @property - def lin_acc_b(self) -> wp.array: + def lin_acc_b(self) -> ProxyArray: """Linear acceleration (coordinate) in the PVA frame [m/s^2]. Equal to -g in freefall, zero at rest. Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). """ - return self._lin_acc_b + if self._lin_acc_b_ta is None: + self._lin_acc_b_ta = ProxyArray(self._lin_acc_b) + return self._lin_acc_b_ta @property - def ang_acc_b(self) -> wp.array: + def ang_acc_b(self) -> ProxyArray: """PVA frame angular acceleration relative to the world expressed in PVA frame. Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). """ - return self._ang_acc_b + if self._ang_acc_b_ta is None: + self._ang_acc_b_ta = ProxyArray(self._ang_acc_b) + return self._ang_acc_b_ta def create_buffers(self, num_envs: int, device: str) -> None: """Create internal buffers for sensor data. @@ -118,3 +135,13 @@ def create_buffers(self, num_envs: int, device: str) -> None: self._ang_vel_b = wp.zeros(num_envs, dtype=wp.vec3f, device=device) self._lin_acc_b = wp.zeros(num_envs, dtype=wp.vec3f, device=device) self._ang_acc_b = wp.zeros(num_envs, dtype=wp.vec3f, device=device) + + # -- Pinned ProxyArray cache (one per read property, lazily created on first access) + self._pose_w_ta: ProxyArray | None = None + self._pos_w_ta: ProxyArray | None = None + self._quat_w_ta: ProxyArray | None = None + self._projected_gravity_b_ta: ProxyArray | None = None + self._lin_vel_b_ta: ProxyArray | None = None + self._ang_vel_b_ta: ProxyArray | None = None + self._lin_acc_b_ta: ProxyArray | None = None + self._ang_acc_b_ta: ProxyArray | None = None diff --git a/source/isaaclab_physx/isaaclab_physx/sim/views/fabric_frame_view.py b/source/isaaclab_physx/isaaclab_physx/sim/views/fabric_frame_view.py index 87adad2238c4..de65e8501793 100644 --- a/source/isaaclab_physx/isaaclab_physx/sim/views/fabric_frame_view.py +++ b/source/isaaclab_physx/isaaclab_physx/sim/views/fabric_frame_view.py @@ -18,6 +18,7 @@ from isaaclab.app.settings_manager import SettingsManager from isaaclab.sim.views.base_frame_view import BaseFrameView from isaaclab.sim.views.usd_frame_view import UsdFrameView +from isaaclab.utils.warp import ProxyArray from isaaclab.utils.warp import fabric as fabric_utils logger = logging.getLogger(__name__) @@ -50,7 +51,7 @@ class FabricFrameView(BaseFrameView): Warp kernels operating on ``omni:fabric:worldMatrix``. All other operations delegate to the internal USD view. - All getters return ``wp.array``. Setters accept ``wp.array``. + Pose getters return :class:`~isaaclab.utils.warp.ProxyArray`. Setters accept ``wp.array``. """ def __init__( @@ -102,6 +103,7 @@ def count(self) -> int: @property def device(self) -> str: + """Device where arrays are allocated (cpu or cuda).""" return self._device @property @@ -168,7 +170,7 @@ def set_world_poses(self, positions=None, orientations=None, indices=None): if self._sync_usd_on_fabric_write: self._usd_view.set_world_poses(positions, orientations, indices) - def get_world_poses(self, indices=None): + def get_world_poses(self, indices: wp.array | None = None) -> tuple[ProxyArray, ProxyArray]: if not self._use_fabric: return self._usd_view.get_world_poses(indices) @@ -204,7 +206,8 @@ def get_world_poses(self, indices=None): if use_cached: wp.synchronize() - return positions_wp, orientations_wp + return self._fabric_positions_ta, self._fabric_orientations_ta + return ProxyArray(positions_wp), ProxyArray(orientations_wp) # ------------------------------------------------------------------ # Local poses — USD fallback (Fabric only accelerates world poses) @@ -213,7 +216,7 @@ def get_world_poses(self, indices=None): def set_local_poses(self, translations=None, orientations=None, indices=None): self._usd_view.set_local_poses(translations, orientations, indices) - def get_local_poses(self, indices=None): + def get_local_poses(self, indices: wp.array | None = None) -> tuple[ProxyArray, ProxyArray]: return self._usd_view.get_local_poses(indices) # ------------------------------------------------------------------ @@ -367,6 +370,8 @@ def _initialize_fabric(self) -> None: self._fabric_positions_buf = wp.zeros((self.count, 3), dtype=wp.float32, device=self._device) self._fabric_orientations_buf = wp.zeros((self.count, 4), dtype=wp.float32, device=self._device) + self._fabric_positions_ta = ProxyArray(self._fabric_positions_buf) + self._fabric_orientations_ta = ProxyArray(self._fabric_orientations_buf) self._fabric_scales_buf = wp.zeros((self.count, 3), dtype=wp.float32, device=self._device) self._fabric_dummy_buffer = wp.zeros((0, 3), dtype=wp.float32, device=self._device) self._fabric_world_matrices = wp.fabricarray(self._fabric_selection, "omni:fabric:worldMatrix") @@ -381,7 +386,9 @@ def _sync_fabric_from_usd_once(self) -> None: if not self._fabric_initialized: self._initialize_fabric() - positions_usd, orientations_usd = self._usd_view.get_world_poses() + positions_usd_ta, orientations_usd_ta = self._usd_view.get_world_poses() + positions_usd = positions_usd_ta.warp + orientations_usd = orientations_usd_ta.warp scales_usd = self._usd_view.get_scales() prev_sync = self._sync_usd_on_fabric_write diff --git a/source/isaaclab_physx/test/assets/test_articulation.py b/source/isaaclab_physx/test/assets/test_articulation.py index cf6f86be0670..6f57999db307 100644 --- a/source/isaaclab_physx/test/assets/test_articulation.py +++ b/source/isaaclab_physx/test/assets/test_articulation.py @@ -233,9 +233,9 @@ def test_initialization_floating_base_non_root(sim, num_articulations, device, a # Check that is fixed base assert not articulation.is_fixed_base # Check buffers that exists and have correct shapes - assert wp.to_torch(articulation.data.root_pos_w).shape == (num_articulations, 3) - assert wp.to_torch(articulation.data.root_quat_w).shape == (num_articulations, 4) - assert wp.to_torch(articulation.data.joint_pos).shape == (num_articulations, 21) + assert articulation.data.root_pos_w.torch.shape == (num_articulations, 3) + assert articulation.data.root_quat_w.torch.shape == (num_articulations, 4) + assert articulation.data.joint_pos.torch.shape == (num_articulations, 21) # Check some internal physx data for debugging # -- joint related @@ -289,11 +289,11 @@ def test_initialization_floating_base(sim, num_articulations, device, add_ground # Check that floating base assert not articulation.is_fixed_base # Check buffers that exists and have correct shapes - assert wp.to_torch(articulation.data.root_pos_w).shape == (num_articulations, 3) - assert wp.to_torch(articulation.data.root_quat_w).shape == (num_articulations, 4) - assert wp.to_torch(articulation.data.joint_pos).shape == (num_articulations, 12) - assert wp.to_torch(articulation.data.body_mass).shape == (num_articulations, articulation.num_bodies) - assert wp.to_torch(articulation.data.body_inertia).shape == (num_articulations, articulation.num_bodies, 9) + assert articulation.data.root_pos_w.torch.shape == (num_articulations, 3) + assert articulation.data.root_quat_w.torch.shape == (num_articulations, 4) + assert articulation.data.joint_pos.torch.shape == (num_articulations, 12) + assert articulation.data.body_mass.torch.shape == (num_articulations, articulation.num_bodies) + assert articulation.data.body_inertia.torch.shape == (num_articulations, articulation.num_bodies, 9) # Check some internal physx data for debugging # -- joint related @@ -346,11 +346,11 @@ def test_initialization_fixed_base(sim, num_articulations, device): # Check that fixed base assert articulation.is_fixed_base # Check buffers that exists and have correct shapes - assert wp.to_torch(articulation.data.root_pos_w).shape == (num_articulations, 3) - assert wp.to_torch(articulation.data.root_quat_w).shape == (num_articulations, 4) - assert wp.to_torch(articulation.data.joint_pos).shape == (num_articulations, 9) - assert wp.to_torch(articulation.data.body_mass).shape == (num_articulations, articulation.num_bodies) - assert wp.to_torch(articulation.data.body_inertia).shape == (num_articulations, articulation.num_bodies, 9) + assert articulation.data.root_pos_w.torch.shape == (num_articulations, 3) + assert articulation.data.root_quat_w.torch.shape == (num_articulations, 4) + assert articulation.data.joint_pos.torch.shape == (num_articulations, 9) + assert articulation.data.body_mass.torch.shape == (num_articulations, articulation.num_bodies) + assert articulation.data.body_inertia.torch.shape == (num_articulations, articulation.num_bodies, 9) # Check some internal physx data for debugging # -- joint related @@ -373,12 +373,12 @@ def test_initialization_fixed_base(sim, num_articulations, device): articulation.update(sim.cfg.dt) # check that the root is at the correct state - its default state as it is fixed base - default_root_pose = wp.to_torch(articulation.data.default_root_pose).clone() - default_root_vel = wp.to_torch(articulation.data.default_root_vel).clone() + default_root_pose = articulation.data.default_root_pose.torch.clone() + default_root_vel = articulation.data.default_root_vel.torch.clone() default_root_pose[:, :3] = default_root_pose[:, :3] + translations - torch.testing.assert_close(wp.to_torch(articulation.data.root_link_pose_w), default_root_pose) - torch.testing.assert_close(wp.to_torch(articulation.data.root_com_vel_w), default_root_vel) + torch.testing.assert_close(articulation.data.root_link_pose_w.torch, default_root_pose) + torch.testing.assert_close(articulation.data.root_com_vel_w.torch, default_root_vel) @pytest.mark.parametrize("num_articulations", [1, 2]) @@ -412,11 +412,11 @@ def test_initialization_fixed_base_single_joint(sim, num_articulations, device, # Check that fixed base assert articulation.is_fixed_base # Check buffers that exists and have correct shapes - assert wp.to_torch(articulation.data.root_pos_w).shape == (num_articulations, 3) - assert wp.to_torch(articulation.data.root_quat_w).shape == (num_articulations, 4) - assert wp.to_torch(articulation.data.joint_pos).shape == (num_articulations, 1) - assert wp.to_torch(articulation.data.body_mass).shape == (num_articulations, articulation.num_bodies) - assert wp.to_torch(articulation.data.body_inertia).shape == (num_articulations, articulation.num_bodies, 9) + assert articulation.data.root_pos_w.torch.shape == (num_articulations, 3) + assert articulation.data.root_quat_w.torch.shape == (num_articulations, 4) + assert articulation.data.joint_pos.torch.shape == (num_articulations, 1) + assert articulation.data.body_mass.torch.shape == (num_articulations, articulation.num_bodies) + assert articulation.data.body_inertia.torch.shape == (num_articulations, articulation.num_bodies, 9) # Check some internal physx data for debugging # -- joint related @@ -439,12 +439,12 @@ def test_initialization_fixed_base_single_joint(sim, num_articulations, device, articulation.update(sim.cfg.dt) # check that the root is at the correct state - its default state as it is fixed base - default_root_pose = wp.to_torch(articulation.data.default_root_pose).clone() - default_root_vel = wp.to_torch(articulation.data.default_root_vel).clone() + default_root_pose = articulation.data.default_root_pose.torch.clone() + default_root_vel = articulation.data.default_root_vel.torch.clone() default_root_pose[:, :3] = default_root_pose[:, :3] + translations - torch.testing.assert_close(wp.to_torch(articulation.data.root_link_pose_w), default_root_pose) - torch.testing.assert_close(wp.to_torch(articulation.data.root_com_vel_w), default_root_vel) + torch.testing.assert_close(articulation.data.root_link_pose_w.torch, default_root_pose) + torch.testing.assert_close(articulation.data.root_com_vel_w.torch, default_root_vel) @pytest.mark.parametrize("num_articulations", [1, 2]) @@ -477,11 +477,11 @@ def test_initialization_hand_with_tendons(sim, num_articulations, device): # Check that fixed base assert articulation.is_fixed_base # Check buffers that exists and have correct shapes - assert wp.to_torch(articulation.data.root_pos_w).shape == (num_articulations, 3) - assert wp.to_torch(articulation.data.root_quat_w).shape == (num_articulations, 4) - assert wp.to_torch(articulation.data.joint_pos).shape == (num_articulations, 24) - assert wp.to_torch(articulation.data.body_mass).shape == (num_articulations, articulation.num_bodies) - assert wp.to_torch(articulation.data.body_inertia).shape == (num_articulations, articulation.num_bodies, 9) + assert articulation.data.root_pos_w.torch.shape == (num_articulations, 3) + assert articulation.data.root_quat_w.torch.shape == (num_articulations, 4) + assert articulation.data.joint_pos.torch.shape == (num_articulations, 24) + assert articulation.data.body_mass.torch.shape == (num_articulations, articulation.num_bodies) + assert articulation.data.body_inertia.torch.shape == (num_articulations, articulation.num_bodies, 9) # Check some internal physx data for debugging # -- joint related @@ -533,9 +533,9 @@ def test_initialization_floating_base_made_fixed_base(sim, num_articulations, de # Check that is fixed base assert articulation.is_fixed_base # Check buffers that exists and have correct shapes - assert wp.to_torch(articulation.data.root_pos_w).shape == (num_articulations, 3) - assert wp.to_torch(articulation.data.root_quat_w).shape == (num_articulations, 4) - assert wp.to_torch(articulation.data.joint_pos).shape == (num_articulations, 12) + assert articulation.data.root_pos_w.torch.shape == (num_articulations, 3) + assert articulation.data.root_quat_w.torch.shape == (num_articulations, 4) + assert articulation.data.joint_pos.torch.shape == (num_articulations, 12) # Check some internal physx data for debugging # -- joint related @@ -554,12 +554,12 @@ def test_initialization_floating_base_made_fixed_base(sim, num_articulations, de articulation.update(sim.cfg.dt) # check that the root is at the correct state - its default state as it is fixed base - default_root_pose = wp.to_torch(articulation.data.default_root_pose).clone() - default_root_vel = wp.to_torch(articulation.data.default_root_vel).clone() + default_root_pose = articulation.data.default_root_pose.torch.clone() + default_root_vel = articulation.data.default_root_vel.torch.clone() default_root_pose[:, :3] = default_root_pose[:, :3] + translations - torch.testing.assert_close(wp.to_torch(articulation.data.root_link_pose_w), default_root_pose) - torch.testing.assert_close(wp.to_torch(articulation.data.root_com_vel_w), default_root_vel) + torch.testing.assert_close(articulation.data.root_link_pose_w.torch, default_root_pose) + torch.testing.assert_close(articulation.data.root_com_vel_w.torch, default_root_vel) @pytest.mark.parametrize("num_articulations", [1, 2]) @@ -594,9 +594,9 @@ def test_initialization_fixed_base_made_floating_base(sim, num_articulations, de # Check that is floating base assert not articulation.is_fixed_base # Check buffers that exists and have correct shapes - assert wp.to_torch(articulation.data.root_pos_w).shape == (num_articulations, 3) - assert wp.to_torch(articulation.data.root_quat_w).shape == (num_articulations, 4) - assert wp.to_torch(articulation.data.joint_pos).shape == (num_articulations, 9) + assert articulation.data.root_pos_w.torch.shape == (num_articulations, 3) + assert articulation.data.root_quat_w.torch.shape == (num_articulations, 4) + assert articulation.data.joint_pos.torch.shape == (num_articulations, 9) # Check some internal physx data for debugging # -- joint related @@ -698,7 +698,7 @@ def test_joint_pos_limits(sim, num_articulations, device, add_ground_plane): assert articulation.is_initialized # Get current default joint pos - default_joint_pos = wp.to_torch(articulation._data.default_joint_pos).clone() + default_joint_pos = articulation._data.default_joint_pos.torch.clone() # Set new joint limits limits = torch.zeros(num_articulations, articulation.num_joints, 2, device=device) @@ -707,8 +707,8 @@ def test_joint_pos_limits(sim, num_articulations, device, add_ground_plane): articulation.write_joint_position_limit_to_sim_index(limits=limits) # Check new limits are in place - torch.testing.assert_close(wp.to_torch(articulation._data.joint_pos_limits), limits) - torch.testing.assert_close(wp.to_torch(articulation._data.default_joint_pos), default_joint_pos) + torch.testing.assert_close(articulation._data.joint_pos_limits.torch, limits) + torch.testing.assert_close(articulation._data.default_joint_pos.torch, default_joint_pos) # Set new joint limits with indexing env_ids = torch.arange(1, device=device, dtype=torch.int32) @@ -719,8 +719,8 @@ def test_joint_pos_limits(sim, num_articulations, device, add_ground_plane): articulation.write_joint_position_limit_to_sim_index(limits=limits, env_ids=env_ids, joint_ids=joint_ids) # Check new limits are in place - torch.testing.assert_close(wp.to_torch(articulation._data.joint_pos_limits)[env_ids][:, joint_ids], limits) - torch.testing.assert_close(wp.to_torch(articulation._data.default_joint_pos), default_joint_pos) + torch.testing.assert_close(articulation._data.joint_pos_limits.torch[env_ids][:, joint_ids], limits) + torch.testing.assert_close(articulation._data.default_joint_pos.torch, default_joint_pos) # Set new joint limits that invalidate default joint pos limits = torch.zeros(num_articulations, articulation.num_joints, 2, device=device) @@ -729,7 +729,7 @@ def test_joint_pos_limits(sim, num_articulations, device, add_ground_plane): articulation.write_joint_position_limit_to_sim_index(limits=limits) # Check if all values are within the bounds - default_joint_pos_torch = wp.to_torch(articulation._data.default_joint_pos) + default_joint_pos_torch = articulation._data.default_joint_pos.torch within_bounds = (default_joint_pos_torch >= limits[..., 0]) & (default_joint_pos_torch <= limits[..., 1]) assert torch.all(within_bounds) @@ -740,7 +740,7 @@ def test_joint_pos_limits(sim, num_articulations, device, add_ground_plane): articulation.write_joint_position_limit_to_sim_index(limits=limits, env_ids=env_ids, joint_ids=joint_ids) # Check if all values are within the bounds - default_joint_pos_torch = wp.to_torch(articulation._data.default_joint_pos) + default_joint_pos_torch = articulation._data.default_joint_pos.torch within_bounds = (default_joint_pos_torch[env_ids][:, joint_ids] >= limits[..., 0]) & ( default_joint_pos_torch[env_ids][:, joint_ids] <= limits[..., 1] ) @@ -768,14 +768,14 @@ def __init__(self, art): assert articulation.is_initialized # Case A: no clipping → should NOT terminate - wp.to_torch(articulation._data.computed_torque).zero_() - wp.to_torch(articulation._data.applied_torque).zero_() + articulation._data.computed_torque.torch.zero_() + articulation._data.applied_torque.torch.zero_() out = joint_effort_out_of_limit(env, robot_all) # [N] assert torch.all(~out) # Case B: simulate clipping → should terminate - wp.to_torch(articulation._data.computed_torque).fill_(100.0) # pretend controller commanded 100 - wp.to_torch(articulation._data.applied_torque).fill_(50.0) # pretend actuator clipped to 50 + articulation._data.computed_torque.torch.fill_(100.0) # pretend controller commanded 100 + articulation._data.applied_torque.torch.fill_(50.0) # pretend actuator clipped to 50 out = joint_effort_out_of_limit(env, robot_all) # [N] assert torch.all(out) @@ -805,13 +805,13 @@ def test_external_force_buffer(sim, num_articulations, device): body_ids, _ = articulation.find_bodies("base") # reset root state - articulation.write_root_pose_to_sim_index(root_pose=wp.to_torch(articulation.data.default_root_pose).clone()) - articulation.write_root_velocity_to_sim_index(root_velocity=wp.to_torch(articulation.data.default_root_vel).clone()) + articulation.write_root_pose_to_sim_index(root_pose=articulation.data.default_root_pose.torch.clone()) + articulation.write_root_velocity_to_sim_index(root_velocity=articulation.data.default_root_vel.torch.clone()) # reset dof state joint_pos, joint_vel = ( - wp.to_torch(articulation.data.default_joint_pos), - wp.to_torch(articulation.data.default_joint_vel), + articulation.data.default_joint_pos.torch, + articulation.data.default_joint_vel.torch, ) articulation.write_joint_position_to_sim_index(position=joint_pos) articulation.write_joint_velocity_to_sim_index(velocity=joint_vel) @@ -844,8 +844,8 @@ def test_external_force_buffer(sim, num_articulations, device): # check if the articulation's force and torque buffers are correctly updated for i in range(num_articulations): - assert wp.to_torch(articulation.permanent_wrench_composer.composed_force)[i, 0, 0].item() == force - assert wp.to_torch(articulation.permanent_wrench_composer.composed_torque)[i, 0, 0].item() == force + assert articulation.permanent_wrench_composer.composed_force.torch[i, 0, 0].item() == force + assert articulation.permanent_wrench_composer.composed_torque.torch[i, 0, 0].item() == force # Check if the instantaneous wrench is correctly added to the permanent wrench articulation.instantaneous_wrench_composer.add_forces_and_torques_index( @@ -855,7 +855,7 @@ def test_external_force_buffer(sim, num_articulations, device): ) # apply action to the articulation - articulation.set_joint_position_target_index(target=wp.to_torch(articulation.data.default_joint_pos).clone()) + articulation.set_joint_position_target_index(target=articulation.data.default_joint_pos.torch.clone()) articulation.write_data_to_sim() # perform step @@ -894,14 +894,12 @@ def test_external_force_on_single_body(sim, num_articulations, device): # Now we are ready! for _ in range(5): # reset root state - articulation.write_root_pose_to_sim_index(root_pose=wp.to_torch(articulation.data.default_root_pose).clone()) - articulation.write_root_velocity_to_sim_index( - root_velocity=wp.to_torch(articulation.data.default_root_vel).clone() - ) + articulation.write_root_pose_to_sim_index(root_pose=articulation.data.default_root_pose.torch.clone()) + articulation.write_root_velocity_to_sim_index(root_velocity=articulation.data.default_root_vel.torch.clone()) # reset dof state joint_pos, joint_vel = ( - wp.to_torch(articulation.data.default_joint_pos), - wp.to_torch(articulation.data.default_joint_vel), + articulation.data.default_joint_pos.torch, + articulation.data.default_joint_vel.torch, ) articulation.write_joint_position_to_sim_index(position=joint_pos) articulation.write_joint_velocity_to_sim_index(velocity=joint_vel) @@ -914,9 +912,7 @@ def test_external_force_on_single_body(sim, num_articulations, device): # perform simulation for _ in range(100): # apply action to the articulation - articulation.set_joint_position_target_index( - target=wp.to_torch(articulation.data.default_joint_pos).clone() - ) + articulation.set_joint_position_target_index(target=articulation.data.default_joint_pos.torch.clone()) articulation.write_data_to_sim() # perform step sim.step() @@ -924,7 +920,7 @@ def test_external_force_on_single_body(sim, num_articulations, device): articulation.update(sim.cfg.dt) # check condition that the articulations have fallen down for i in range(num_articulations): - assert wp.to_torch(articulation.data.root_pos_w)[i, 2].item() < 0.2 + assert articulation.data.root_pos_w.torch[i, 2].item() < 0.2 @pytest.mark.parametrize("num_articulations", [1, 2]) @@ -965,17 +961,15 @@ def test_external_force_on_single_body_at_position(sim, num_articulations, devic # Now we are ready! for i in range(5): # reset root state - root_pose = wp.to_torch(articulation.data.default_root_pose).clone() + root_pose = articulation.data.default_root_pose.torch.clone() root_pose[0, 0] = 2.5 # space them apart by 2.5m articulation.write_root_pose_to_sim_index(root_pose=root_pose) - articulation.write_root_velocity_to_sim_index( - root_velocity=wp.to_torch(articulation.data.default_root_vel).clone() - ) + articulation.write_root_velocity_to_sim_index(root_velocity=articulation.data.default_root_vel.torch.clone()) # reset dof state joint_pos, joint_vel = ( - wp.to_torch(articulation.data.default_joint_pos), - wp.to_torch(articulation.data.default_joint_vel), + articulation.data.default_joint_pos.torch, + articulation.data.default_joint_vel.torch, ) articulation.write_joint_position_to_sim_index(position=joint_pos) articulation.write_joint_velocity_to_sim_index(velocity=joint_vel) @@ -985,7 +979,7 @@ def test_external_force_on_single_body_at_position(sim, num_articulations, devic is_global = False if i % 2 == 0: - body_com_pos_w = wp.to_torch(articulation.data.body_com_pos_w)[:, body_ids, :3] + body_com_pos_w = articulation.data.body_com_pos_w.torch[:, body_ids, :3] # is_global = True external_wrench_positions_b[..., 0] = 0.0 external_wrench_positions_b[..., 1] = 1.0 @@ -1013,9 +1007,7 @@ def test_external_force_on_single_body_at_position(sim, num_articulations, devic # perform simulation for _ in range(100): # apply action to the articulation - articulation.set_joint_position_target_index( - target=wp.to_torch(articulation.data.default_joint_pos).clone() - ) + articulation.set_joint_position_target_index(target=articulation.data.default_joint_pos.torch.clone()) articulation.write_data_to_sim() # perform step sim.step() @@ -1023,7 +1015,7 @@ def test_external_force_on_single_body_at_position(sim, num_articulations, devic articulation.update(sim.cfg.dt) # check condition that the articulations have fallen down for i in range(num_articulations): - assert wp.to_torch(articulation.data.root_pos_w)[i, 2].item() < 0.2 + assert articulation.data.root_pos_w.torch[i, 2].item() < 0.2 @pytest.mark.parametrize("num_articulations", [1, 2]) @@ -1056,14 +1048,12 @@ def test_external_force_on_multiple_bodies(sim, num_articulations, device): # Now we are ready! for _ in range(5): # reset root state - articulation.write_root_pose_to_sim_index(root_pose=wp.to_torch(articulation.data.default_root_pose).clone()) - articulation.write_root_velocity_to_sim_index( - root_velocity=wp.to_torch(articulation.data.default_root_vel).clone() - ) + articulation.write_root_pose_to_sim_index(root_pose=articulation.data.default_root_pose.torch.clone()) + articulation.write_root_velocity_to_sim_index(root_velocity=articulation.data.default_root_vel.torch.clone()) # reset dof state joint_pos, joint_vel = ( - wp.to_torch(articulation.data.default_joint_pos), - wp.to_torch(articulation.data.default_joint_vel), + articulation.data.default_joint_pos.torch, + articulation.data.default_joint_vel.torch, ) articulation.write_joint_position_to_sim_index(position=joint_pos) articulation.write_joint_velocity_to_sim_index(velocity=joint_vel) @@ -1076,9 +1066,7 @@ def test_external_force_on_multiple_bodies(sim, num_articulations, device): # perform simulation for _ in range(100): # apply action to the articulation - articulation.set_joint_position_target_index( - target=wp.to_torch(articulation.data.default_joint_pos).clone() - ) + articulation.set_joint_position_target_index(target=articulation.data.default_joint_pos.torch.clone()) articulation.write_data_to_sim() # perform step sim.step() @@ -1087,7 +1075,7 @@ def test_external_force_on_multiple_bodies(sim, num_articulations, device): # check condition for i in range(num_articulations): # since there is a moment applied on the articulation, the articulation should rotate - assert wp.to_torch(articulation.data.root_ang_vel_w)[i, 2].item() > 0.1 + assert articulation.data.root_ang_vel_w.torch[i, 2].item() > 0.1 @pytest.mark.parametrize("num_articulations", [1, 2]) @@ -1129,14 +1117,12 @@ def test_external_force_on_multiple_bodies_at_position(sim, num_articulations, d # Now we are ready! for i in range(5): # reset root state - articulation.write_root_pose_to_sim_index(root_pose=wp.to_torch(articulation.data.default_root_pose).clone()) - articulation.write_root_velocity_to_sim_index( - root_velocity=wp.to_torch(articulation.data.default_root_vel).clone() - ) + articulation.write_root_pose_to_sim_index(root_pose=articulation.data.default_root_pose.torch.clone()) + articulation.write_root_velocity_to_sim_index(root_velocity=articulation.data.default_root_vel.torch.clone()) # reset dof state joint_pos, joint_vel = ( - wp.to_torch(articulation.data.default_joint_pos), - wp.to_torch(articulation.data.default_joint_vel), + articulation.data.default_joint_pos.torch, + articulation.data.default_joint_vel.torch, ) articulation.write_joint_position_to_sim_index(position=joint_pos) articulation.write_joint_velocity_to_sim_index(velocity=joint_vel) @@ -1145,7 +1131,7 @@ def test_external_force_on_multiple_bodies_at_position(sim, num_articulations, d is_global = False if i % 2 == 0: - body_com_pos_w = wp.to_torch(articulation.data.body_com_pos_w)[:, body_ids, :3] + body_com_pos_w = articulation.data.body_com_pos_w.torch[:, body_ids, :3] is_global = True external_wrench_positions_b[..., 0] = 0.0 external_wrench_positions_b[..., 1] = 1.0 @@ -1174,9 +1160,7 @@ def test_external_force_on_multiple_bodies_at_position(sim, num_articulations, d # perform simulation for _ in range(100): # apply action to the articulation - articulation.set_joint_position_target_index( - target=wp.to_torch(articulation.data.default_joint_pos).clone() - ) + articulation.set_joint_position_target_index(target=articulation.data.default_joint_pos.torch.clone()) articulation.write_data_to_sim() # perform step sim.step() @@ -1185,7 +1169,7 @@ def test_external_force_on_multiple_bodies_at_position(sim, num_articulations, d # check condition for i in range(num_articulations): # since there is a moment applied on the articulation, the articulation should rotate - assert torch.abs(wp.to_torch(articulation.data.root_ang_vel_w)[i, 2]).item() > 0.1 + assert torch.abs(articulation.data.root_ang_vel_w.torch[i, 2]).item() > 0.1 @pytest.mark.parametrize("num_articulations", [1, 2]) @@ -1360,7 +1344,7 @@ def test_setting_velocity_limit_implicit(sim, num_articulations, device, vel_lim # read the values set into the simulation physx_vel_limit = wp.to_torch(articulation.root_view.get_dof_max_velocities()).to(device) # check data buffer - torch.testing.assert_close(wp.to_torch(articulation.data.joint_velocity_limits), physx_vel_limit) + torch.testing.assert_close(articulation.data.joint_velocity_limits.torch, physx_vel_limit) # check actuator has simulation velocity limit torch.testing.assert_close(articulation.actuators["joint"].velocity_limit_sim, physx_vel_limit) # check that both values match for velocity limit @@ -1412,7 +1396,7 @@ def test_setting_velocity_limit_explicit(sim, num_articulations, device, vel_lim actuator_vel_limit_sim = articulation.actuators["joint"].velocity_limit_sim # check data buffer for joint_velocity_limits_sim - torch.testing.assert_close(wp.to_torch(articulation.data.joint_velocity_limits), physx_vel_limit) + torch.testing.assert_close(articulation.data.joint_velocity_limits.torch, physx_vel_limit) # check actuator velocity_limit_sim is set to physx torch.testing.assert_close(actuator_vel_limit_sim, physx_vel_limit) @@ -1575,10 +1559,10 @@ def test_reset(sim, num_articulations, device): # Reset should zero external forces and torques assert not articulation._instantaneous_wrench_composer.active assert not articulation._permanent_wrench_composer.active - assert torch.count_nonzero(wp.to_torch(articulation._instantaneous_wrench_composer.composed_force)) == 0 - assert torch.count_nonzero(wp.to_torch(articulation._instantaneous_wrench_composer.composed_torque)) == 0 - assert torch.count_nonzero(wp.to_torch(articulation._permanent_wrench_composer.composed_force)) == 0 - assert torch.count_nonzero(wp.to_torch(articulation._permanent_wrench_composer.composed_torque)) == 0 + assert torch.count_nonzero(articulation._instantaneous_wrench_composer.composed_force.torch) == 0 + assert torch.count_nonzero(articulation._instantaneous_wrench_composer.composed_torque.torch) == 0 + assert torch.count_nonzero(articulation._permanent_wrench_composer.composed_force.torch) == 0 + assert torch.count_nonzero(articulation._permanent_wrench_composer.composed_torque.torch) == 0 if num_articulations > 1: num_bodies = articulation.num_bodies @@ -1593,20 +1577,10 @@ def test_reset(sim, num_articulations, device): articulation.reset(env_ids=torch.tensor([0], device=device)) assert articulation._instantaneous_wrench_composer.active assert articulation._permanent_wrench_composer.active - assert ( - torch.count_nonzero(wp.to_torch(articulation._instantaneous_wrench_composer.composed_force)) - == num_bodies * 3 - ) - assert ( - torch.count_nonzero(wp.to_torch(articulation._instantaneous_wrench_composer.composed_torque)) - == num_bodies * 3 - ) - assert ( - torch.count_nonzero(wp.to_torch(articulation._permanent_wrench_composer.composed_force)) == num_bodies * 3 - ) - assert ( - torch.count_nonzero(wp.to_torch(articulation._permanent_wrench_composer.composed_torque)) == num_bodies * 3 - ) + assert torch.count_nonzero(articulation._instantaneous_wrench_composer.composed_force.torch) == num_bodies * 3 + assert torch.count_nonzero(articulation._instantaneous_wrench_composer.composed_torque.torch) == num_bodies * 3 + assert torch.count_nonzero(articulation._permanent_wrench_composer.composed_force.torch) == num_bodies * 3 + assert torch.count_nonzero(articulation._permanent_wrench_composer.composed_torque.torch) == num_bodies * 3 @pytest.mark.parametrize("num_articulations", [1, 2]) @@ -1630,7 +1604,7 @@ def test_apply_joint_command(sim, num_articulations, device, add_ground_plane): articulation.update(sim.cfg.dt) # reset dof state - joint_pos = wp.to_torch(articulation.data.default_joint_pos).clone() + joint_pos = articulation.data.default_joint_pos.torch.clone() joint_pos[:, 3] = 0.0 # apply action to the articulation @@ -1646,7 +1620,7 @@ def test_apply_joint_command(sim, num_articulations, device, add_ground_plane): # Check that current joint position is not the same as default joint position, meaning # the articulation moved. We can't check that it reached its desired joint position as the gains # are not properly tuned - assert not torch.allclose(wp.to_torch(articulation.data.joint_pos), joint_pos) + assert not torch.allclose(articulation.data.joint_pos.torch, joint_pos) @pytest.mark.parametrize("num_articulations", [1, 2]) @@ -1710,19 +1684,19 @@ def test_body_root_state(sim, num_articulations, device, with_offset): articulation.update(sim.cfg.dt) # get state properties - root_link_pose_w = wp.to_torch(articulation.data.root_link_pose_w) - root_link_vel_w = wp.to_torch(articulation.data.root_link_vel_w) - root_com_pose_w = wp.to_torch(articulation.data.root_com_pose_w) - root_com_vel_w = wp.to_torch(articulation.data.root_com_vel_w) - body_link_pose_w = wp.to_torch(articulation.data.body_link_pose_w) - body_link_vel_w = wp.to_torch(articulation.data.body_link_vel_w) - body_com_pose_w = wp.to_torch(articulation.data.body_com_pose_w) - body_com_vel_w = wp.to_torch(articulation.data.body_com_vel_w) + root_link_pose_w = articulation.data.root_link_pose_w.torch + root_link_vel_w = articulation.data.root_link_vel_w.torch + root_com_pose_w = articulation.data.root_com_pose_w.torch + root_com_vel_w = articulation.data.root_com_vel_w.torch + body_link_pose_w = articulation.data.body_link_pose_w.torch + body_link_vel_w = articulation.data.body_link_vel_w.torch + body_com_pose_w = articulation.data.body_com_pose_w.torch + body_com_vel_w = articulation.data.body_com_vel_w.torch if with_offset: # get joint state - joint_pos = wp.to_torch(articulation.data.joint_pos).unsqueeze(-1) - joint_vel = wp.to_torch(articulation.data.joint_vel).unsqueeze(-1) + joint_pos = articulation.data.joint_pos.torch.unsqueeze(-1) + joint_vel = articulation.data.joint_vel.torch.unsqueeze(-1) # LINK state # angular velocity should be the same for both COM and link frames @@ -1757,7 +1731,7 @@ def test_body_root_state(sim, num_articulations, device, with_offset): torch.testing.assert_close(pos_gt, body_com_pose_w[..., :3], atol=1e-3, rtol=1e-1) # orientation - com_quat_b = wp.to_torch(articulation.data.body_com_quat_b) + com_quat_b = articulation.data.body_com_quat_b.torch com_quat_w = math_utils.quat_mul(body_link_pose_w[..., 3:], com_quat_b) torch.testing.assert_close(com_quat_w, body_com_pose_w[..., 3:]) torch.testing.assert_close(com_quat_w[:, root_idx, :], root_com_pose_w[..., 3:]) @@ -1821,7 +1795,7 @@ def test_write_root_state(sim, num_articulations, device, with_offset, state_loc torch.testing.assert_close(wp.to_torch(articulation.root_view.get_coms()), com) rand_state = torch.zeros(num_articulations, 13, device=device) - rand_state[..., :7] = wp.to_torch(articulation.data.default_root_pose) + rand_state[..., :7] = articulation.data.default_root_pose.torch rand_state[..., :3] += env_pos # make quaternion a unit vector rand_state[..., 3:7] = torch.nn.functional.normalize(rand_state[..., 3:7], dim=-1) @@ -1849,11 +1823,11 @@ def test_write_root_state(sim, num_articulations, device, with_offset, state_loc articulation.write_root_link_velocity_to_sim_index(root_velocity=rand_state[..., 7:], env_ids=env_idx) if state_location == "com": - torch.testing.assert_close(rand_state[..., :7], wp.to_torch(articulation.data.root_com_pose_w)) - torch.testing.assert_close(rand_state[..., 7:], wp.to_torch(articulation.data.root_com_vel_w)) + torch.testing.assert_close(rand_state[..., :7], articulation.data.root_com_pose_w.torch) + torch.testing.assert_close(rand_state[..., 7:], articulation.data.root_com_vel_w.torch) elif state_location == "link": - torch.testing.assert_close(rand_state[..., :7], wp.to_torch(articulation.data.root_link_pose_w)) - torch.testing.assert_close(rand_state[..., 7:], wp.to_torch(articulation.data.root_link_vel_w)) + torch.testing.assert_close(rand_state[..., :7], articulation.data.root_link_pose_w.torch) + torch.testing.assert_close(rand_state[..., 7:], articulation.data.root_link_vel_w.torch) @pytest.mark.parametrize("num_articulations", [1, 2]) @@ -1890,12 +1864,12 @@ def test_body_incoming_joint_wrench_b_single_joint(sim, num_articulations, devic external_torque_vector_b[:, arm_idx, 2] = 10.0 # 10 Nm in z direction # apply action to the articulation - joint_pos = torch.ones_like(wp.to_torch(articulation.data.joint_pos)) * 1.5708 / 2.0 + joint_pos = torch.ones_like(articulation.data.joint_pos.torch) * 1.5708 / 2.0 articulation.write_joint_position_to_sim_index( - position=torch.ones_like(wp.to_torch(articulation.data.joint_pos)), + position=torch.ones_like(articulation.data.joint_pos.torch), ) articulation.write_joint_velocity_to_sim_index( - velocity=torch.zeros_like(wp.to_torch(articulation.data.joint_vel)), + velocity=torch.zeros_like(articulation.data.joint_vel.torch), ) articulation.set_joint_position_target_index(target=joint_pos) articulation.write_data_to_sim() @@ -1910,16 +1884,16 @@ def test_body_incoming_joint_wrench_b_single_joint(sim, num_articulations, devic articulation.update(sim.cfg.dt) # check shape - assert wp.to_torch(articulation.data.body_incoming_joint_wrench_b).shape == ( + assert articulation.data.body_incoming_joint_wrench_b.torch.shape == ( num_articulations, articulation.num_bodies, 6, ) # calculate expected static - mass = wp.to_torch(articulation.data.body_mass).to("cpu") - pos_w = wp.to_torch(articulation.data.body_pos_w) - quat_w = wp.to_torch(articulation.data.body_quat_w) + mass = articulation.data.body_mass.torch.to("cpu") + pos_w = articulation.data.body_pos_w.torch + quat_w = articulation.data.body_quat_w.torch mass_link2 = mass[:, arm_idx].view(num_articulations, -1) gravity = torch.tensor(sim.cfg.gravity, device="cpu").repeat(num_articulations, 1).view((num_articulations, 3)) @@ -1952,7 +1926,7 @@ def test_body_incoming_joint_wrench_b_single_joint(sim, num_articulations, devic # check value of last joint wrench torch.testing.assert_close( expected_wrench, - wp.to_torch(articulation.data.body_incoming_joint_wrench_b)[:, arm_idx, :].squeeze(1), + articulation.data.body_incoming_joint_wrench_b.torch[:, arm_idx, :].squeeze(1), atol=1e-2, rtol=1e-3, ) @@ -2025,13 +1999,13 @@ def test_write_joint_state_data_consistency(sim, num_articulations, device, grav from torch.distributions import Uniform - joint_pos_limits = wp.to_torch(articulation.data.joint_pos_limits) - joint_vel_limits = wp.to_torch(articulation.data.joint_vel_limits) + joint_pos_limits = articulation.data.joint_pos_limits.torch + joint_vel_limits = articulation.data.joint_vel_limits.torch pos_dist = Uniform(joint_pos_limits[..., 0], joint_pos_limits[..., 1]) vel_dist = Uniform(-joint_vel_limits, joint_vel_limits) - original_body_link_pose_w = wp.to_torch(articulation.data.body_link_pose_w).clone() - original_body_com_vel_w = wp.to_torch(articulation.data.body_com_vel_w).clone() + original_body_link_pose_w = articulation.data.body_link_pose_w.torch.clone() + original_body_com_vel_w = articulation.data.body_com_vel_w.torch.clone() rand_joint_pos = pos_dist.sample() rand_joint_vel = vel_dist.sample() @@ -2039,65 +2013,65 @@ def test_write_joint_state_data_consistency(sim, num_articulations, device, grav articulation.write_joint_position_to_sim_index(position=rand_joint_pos) articulation.write_joint_velocity_to_sim_index(velocity=rand_joint_vel) # make sure valued updated - body_link_pose_w = wp.to_torch(articulation.data.body_link_pose_w) - body_com_vel_w = wp.to_torch(articulation.data.body_com_vel_w) + body_link_pose_w = articulation.data.body_link_pose_w.torch + body_com_vel_w = articulation.data.body_com_vel_w.torch original_body_states = torch.cat([original_body_link_pose_w, original_body_com_vel_w], dim=-1) body_state_w = torch.cat([body_link_pose_w, body_com_vel_w], dim=-1) assert torch.count_nonzero(original_body_states[:, 1:] != body_state_w[:, 1:]) > ( len(original_body_states[:, 1:]) / 2 ) # validate body - link consistency - body_link_vel_w = wp.to_torch(articulation.data.body_link_vel_w) - torch.testing.assert_close(body_link_pose_w, wp.to_torch(articulation.data.body_link_pose_w)) + body_link_vel_w = articulation.data.body_link_vel_w.torch + torch.testing.assert_close(body_link_pose_w, articulation.data.body_link_pose_w.torch) # skip lin_vel because it differs from link frame, this should be fine because we are only checking # if velocity update is triggered, which can be determined by comparing angular velocity torch.testing.assert_close(body_com_vel_w[..., 3:], body_link_vel_w[..., 3:]) # validate link - com conistency - body_com_pos_b = wp.to_torch(articulation.data.body_com_pos_b) - body_com_quat_b = wp.to_torch(articulation.data.body_com_quat_b) + body_com_pos_b = articulation.data.body_com_pos_b.torch + body_com_quat_b = articulation.data.body_com_quat_b.torch expected_com_pos, expected_com_quat = math_utils.combine_frame_transforms( body_link_pose_w[..., :3].view(-1, 3), body_link_pose_w[..., 3:].view(-1, 4), body_com_pos_b.view(-1, 3), body_com_quat_b.view(-1, 4), ) - body_com_pos_w = wp.to_torch(articulation.data.body_com_pos_w) - body_com_quat_w = wp.to_torch(articulation.data.body_com_quat_w) + body_com_pos_w = articulation.data.body_com_pos_w.torch + body_com_quat_w = articulation.data.body_com_quat_w.torch torch.testing.assert_close(expected_com_pos.view(len(env_idx), -1, 3), body_com_pos_w) torch.testing.assert_close(expected_com_quat.view(len(env_idx), -1, 4), body_com_quat_w) # validate body - com consistency - body_com_lin_vel_w = wp.to_torch(articulation.data.body_com_lin_vel_w) - body_com_ang_vel_w = wp.to_torch(articulation.data.body_com_ang_vel_w) + body_com_lin_vel_w = articulation.data.body_com_lin_vel_w.torch + body_com_ang_vel_w = articulation.data.body_com_ang_vel_w.torch torch.testing.assert_close(body_com_vel_w[..., :3], body_com_lin_vel_w) torch.testing.assert_close(body_com_vel_w[..., 3:], body_com_ang_vel_w) # validate pos_w, quat_w, pos_b, quat_b is consistent with pose_w and pose_b expected_com_pose_w = torch.cat((body_com_pos_w, body_com_quat_w), dim=2) expected_com_pose_b = torch.cat((body_com_pos_b, body_com_quat_b), dim=2) - body_pos_w = wp.to_torch(articulation.data.body_pos_w) - body_quat_w = wp.to_torch(articulation.data.body_quat_w) + body_pos_w = articulation.data.body_pos_w.torch + body_quat_w = articulation.data.body_quat_w.torch expected_body_pose_w = torch.cat((body_pos_w, body_quat_w), dim=2) - body_link_pos_w = wp.to_torch(articulation.data.body_link_pos_w) - body_link_quat_w = wp.to_torch(articulation.data.body_link_quat_w) + body_link_pos_w = articulation.data.body_link_pos_w.torch + body_link_quat_w = articulation.data.body_link_quat_w.torch expected_body_link_pose_w = torch.cat((body_link_pos_w, body_link_quat_w), dim=2) - body_com_pose_w = wp.to_torch(articulation.data.body_com_pose_w) - body_com_pose_b = wp.to_torch(articulation.data.body_com_pose_b) - body_pose_w = wp.to_torch(articulation.data.body_pose_w) - body_link_pose_w_fresh = wp.to_torch(articulation.data.body_link_pose_w) + body_com_pose_w = articulation.data.body_com_pose_w.torch + body_com_pose_b = articulation.data.body_com_pose_b.torch + body_pose_w = articulation.data.body_pose_w.torch + body_link_pose_w_fresh = articulation.data.body_link_pose_w.torch torch.testing.assert_close(body_com_pose_w, expected_com_pose_w) torch.testing.assert_close(body_com_pose_b, expected_com_pose_b) torch.testing.assert_close(body_pose_w, expected_body_pose_w) torch.testing.assert_close(body_link_pose_w_fresh, expected_body_link_pose_w) # validate pose_w is consistent with individual properties - body_vel_w = wp.to_torch(articulation.data.body_vel_w) - body_com_vel_w_fresh = wp.to_torch(articulation.data.body_com_vel_w) + body_vel_w = articulation.data.body_vel_w.torch + body_com_vel_w_fresh = articulation.data.body_com_vel_w.torch torch.testing.assert_close(body_pose_w, body_link_pose_w) torch.testing.assert_close(body_vel_w, body_com_vel_w) torch.testing.assert_close(body_link_pose_w_fresh, body_link_pose_w) - torch.testing.assert_close(body_com_pose_w, wp.to_torch(articulation.data.body_com_pose_w)) + torch.testing.assert_close(body_com_pose_w, articulation.data.body_com_pose_w.torch) torch.testing.assert_close(body_vel_w, body_com_vel_w_fresh) @@ -2132,11 +2106,11 @@ def test_spatial_tendons(sim, num_articulations, device): # Check that fixed base assert articulation.is_fixed_base # Check buffers that exists and have correct shapes - assert wp.to_torch(articulation.data.root_pos_w).shape == (num_articulations, 3) - assert wp.to_torch(articulation.data.root_quat_w).shape == (num_articulations, 4) - assert wp.to_torch(articulation.data.joint_pos).shape == (num_articulations, 3) - assert wp.to_torch(articulation.data.body_mass).shape == (num_articulations, articulation.num_bodies) - assert wp.to_torch(articulation.data.body_inertia).shape == (num_articulations, articulation.num_bodies, 9) + assert articulation.data.root_pos_w.torch.shape == (num_articulations, 3) + assert articulation.data.root_quat_w.torch.shape == (num_articulations, 4) + assert articulation.data.joint_pos.torch.shape == (num_articulations, 3) + assert articulation.data.body_mass.torch.shape == (num_articulations, articulation.num_bodies) + assert articulation.data.body_inertia.torch.shape == (num_articulations, articulation.num_bodies, 9) assert articulation.num_spatial_tendons == 1 articulation.set_spatial_tendon_stiffness_index(stiffness=10.0) diff --git a/source/isaaclab_physx/test/assets/test_deformable_object.py b/source/isaaclab_physx/test/assets/test_deformable_object.py index 92bd9d145e2d..31a199938c94 100644 --- a/source/isaaclab_physx/test/assets/test_deformable_object.py +++ b/source/isaaclab_physx/test/assets/test_deformable_object.py @@ -141,16 +141,16 @@ def test_initialization(sim, num_cubes, material_path): # Check buffers that exist and have correct shapes # nodal_state_w is (N, V) vec6f -> wp.to_torch gives (N, V, 6) - assert wp.to_torch(cube_object.data.nodal_state_w).shape == (num_cubes, cube_object.max_sim_vertices_per_body, 6) - # nodal_kinematic_target is (N, V) vec4f -> wp.to_torch gives (N, V, 4) - assert wp.to_torch(cube_object.data.nodal_kinematic_target).shape == ( + assert cube_object.data.nodal_state_w.torch.shape == (num_cubes, cube_object.max_sim_vertices_per_body, 6) + # nodal_kinematic_target is (N, V) vec4f -> .torch gives (N, V, 4) + assert cube_object.data.nodal_kinematic_target.torch.shape == ( num_cubes, cube_object.max_sim_vertices_per_body, 4, ) # root_pos_w is (N,) vec3f -> wp.to_torch gives (N, 3) - assert wp.to_torch(cube_object.data.root_pos_w).shape == (num_cubes, 3) - assert wp.to_torch(cube_object.data.root_vel_w).shape == (num_cubes, 3) + assert cube_object.data.root_pos_w.torch.shape == (num_cubes, 3) + assert cube_object.data.root_vel_w.torch.shape == (num_cubes, 3) @pytest.mark.parametrize("num_cubes", [1, 2]) @@ -175,9 +175,9 @@ def test_initialization_surface_deformable(sim, num_cubes): assert cube_object.material_physx_view.count == num_cubes # Check nodal state buffers have correct shapes - assert wp.to_torch(cube_object.data.nodal_state_w).shape == (num_cubes, cube_object.max_sim_vertices_per_body, 6) - assert wp.to_torch(cube_object.data.root_pos_w).shape == (num_cubes, 3) - assert wp.to_torch(cube_object.data.root_vel_w).shape == (num_cubes, 3) + assert cube_object.data.nodal_state_w.torch.shape == (num_cubes, cube_object.max_sim_vertices_per_body, 6) + assert cube_object.data.root_pos_w.torch.shape == (num_cubes, 3) + assert cube_object.data.root_vel_w.torch.shape == (num_cubes, 3) # Kinematic targets are not allocated for surface deformables assert cube_object.data.nodal_kinematic_target is None @@ -214,8 +214,8 @@ def test_set_nodal_state(sim, num_cubes): for state_type_to_randomize in ["nodal_pos_w", "nodal_vel_w"]: state_dict = { - "nodal_pos_w": torch.zeros_like(wp.to_torch(cube_object.data.nodal_pos_w)), - "nodal_vel_w": torch.zeros_like(wp.to_torch(cube_object.data.nodal_vel_w)), + "nodal_pos_w": torch.zeros_like(cube_object.data.nodal_pos_w.torch), + "nodal_vel_w": torch.zeros_like(cube_object.data.nodal_vel_w.torch), } for _ in range(5): @@ -235,9 +235,7 @@ def test_set_nodal_state(sim, num_cubes): ) cube_object.write_nodal_state_to_sim_index(nodal_state) - torch.testing.assert_close( - wp.to_torch(cube_object.data.nodal_state_w), nodal_state, rtol=1e-5, atol=1e-5 - ) + torch.testing.assert_close(cube_object.data.nodal_state_w.torch, nodal_state, rtol=1e-5, atol=1e-5) sim.step() cube_object.update(sim.cfg.dt) @@ -260,7 +258,7 @@ def test_set_nodal_state_with_applied_transform(num_cubes, randomize_pos, random sim.reset() for _ in range(5): - nodal_state = wp.to_torch(cube_object.data.default_nodal_state_w).clone() + nodal_state = cube_object.data.default_nodal_state_w.torch.clone() mean_nodal_pos_default = nodal_state[..., :3].mean(dim=1) if randomize_pos: @@ -288,9 +286,7 @@ def test_set_nodal_state_with_applied_transform(num_cubes, randomize_pos, random sim.step() cube_object.update(sim.cfg.dt) - torch.testing.assert_close( - wp.to_torch(cube_object.data.root_pos_w), mean_nodal_pos_init, rtol=1e-4, atol=1e-4 - ) + torch.testing.assert_close(cube_object.data.root_pos_w.torch, mean_nodal_pos_init, rtol=1e-4, atol=1e-4) @pytest.mark.parametrize("num_cubes", [2, 4]) @@ -304,15 +300,15 @@ def test_set_kinematic_targets(sim, num_cubes): nodal_kinematic_targets = wp.to_torch(cube_object.root_view.get_simulation_nodal_kinematic_targets()).clone() for _ in range(5): - cube_object.write_nodal_state_to_sim_index(wp.to_torch(cube_object.data.default_nodal_state_w)) + cube_object.write_nodal_state_to_sim_index(cube_object.data.default_nodal_state_w.torch) - default_root_pos = wp.to_torch(cube_object.data.default_nodal_state_w).mean(dim=1) + default_root_pos = cube_object.data.default_nodal_state_w.torch.mean(dim=1) cube_object.reset() nodal_kinematic_targets[1:, :, 3] = 1.0 nodal_kinematic_targets[0, :, 3] = 0.0 - nodal_kinematic_targets[0, :, :3] = wp.to_torch(cube_object.data.default_nodal_state_w)[0, :, :3] + nodal_kinematic_targets[0, :, :3] = cube_object.data.default_nodal_state_w.torch[0, :, :3] cube_object.write_nodal_kinematic_target_to_sim_index( nodal_kinematic_targets[0:1], env_ids=torch.tensor([0], device=sim.device) ) @@ -322,10 +318,10 @@ def test_set_kinematic_targets(sim, num_cubes): cube_object.update(sim.cfg.dt) torch.testing.assert_close( - wp.to_torch(cube_object.data.nodal_pos_w)[0], + cube_object.data.nodal_pos_w.torch[0], nodal_kinematic_targets[0, :, :3], rtol=1e-5, atol=1e-5, ) - root_pos_w = wp.to_torch(cube_object.data.root_pos_w) + root_pos_w = cube_object.data.root_pos_w.torch assert torch.all(root_pos_w[1:, 2] < default_root_pos[1:, 2]) diff --git a/source/isaaclab_physx/test/assets/test_rigid_object.py b/source/isaaclab_physx/test/assets/test_rigid_object.py index a41f5a39f083..6a2211787e19 100644 --- a/source/isaaclab_physx/test/assets/test_rigid_object.py +++ b/source/isaaclab_physx/test/assets/test_rigid_object.py @@ -118,10 +118,10 @@ def test_initialization(num_cubes, device): assert len(cube_object.body_names) == 1 # Check buffers that exists and have correct shapes - assert wp.to_torch(cube_object.data.root_pos_w).shape == (num_cubes, 3) - assert wp.to_torch(cube_object.data.root_quat_w).shape == (num_cubes, 4) - assert wp.to_torch(cube_object.data.body_mass).shape == (num_cubes, 1) - assert wp.to_torch(cube_object.data.body_inertia).shape == (num_cubes, 1, 9) + assert cube_object.data.root_pos_w.torch.shape == (num_cubes, 3) + assert cube_object.data.root_quat_w.torch.shape == (num_cubes, 4) + assert cube_object.data.body_mass.torch.shape == (num_cubes, 1) + assert cube_object.data.body_inertia.torch.shape == (num_cubes, 1, 9) # Simulate physics for _ in range(2): @@ -152,8 +152,8 @@ def test_initialization_with_kinematic_enabled(num_cubes, device): assert len(cube_object.body_names) == 1 # Check buffers that exists and have correct shapes - assert wp.to_torch(cube_object.data.root_pos_w).shape == (num_cubes, 3) - assert wp.to_torch(cube_object.data.root_quat_w).shape == (num_cubes, 4) + assert cube_object.data.root_pos_w.torch.shape == (num_cubes, 3) + assert cube_object.data.root_quat_w.torch.shape == (num_cubes, 4) # Simulate physics for _ in range(2): @@ -162,11 +162,11 @@ def test_initialization_with_kinematic_enabled(num_cubes, device): # update object cube_object.update(sim.cfg.dt) # check that the object is kinematic - default_root_pose = wp.to_torch(cube_object.data.default_root_pose).clone() - default_root_vel = wp.to_torch(cube_object.data.default_root_vel).clone() + default_root_pose = cube_object.data.default_root_pose.torch.clone() + default_root_vel = cube_object.data.default_root_vel.torch.clone() default_root_pose[:, :3] += origins - torch.testing.assert_close(wp.to_torch(cube_object.data.root_link_pose_w), default_root_pose) - torch.testing.assert_close(wp.to_torch(cube_object.data.root_com_vel_w), default_root_vel) + torch.testing.assert_close(cube_object.data.root_link_pose_w.torch, default_root_pose) + torch.testing.assert_close(cube_object.data.root_com_vel_w.torch, default_root_vel) @pytest.mark.parametrize("num_cubes", [1, 2]) @@ -253,8 +253,8 @@ def test_external_force_buffer(device): # check if the cube's force and torque buffers are correctly updated for i in range(cube_object.num_instances): - assert wp.to_torch(cube_object._permanent_wrench_composer.composed_force)[i, 0, 0].item() == force - assert wp.to_torch(cube_object._permanent_wrench_composer.composed_torque)[i, 0, 0].item() == force + assert cube_object._permanent_wrench_composer.composed_force.torch[i, 0, 0].item() == force + assert cube_object._permanent_wrench_composer.composed_torque.torch[i, 0, 0].item() == force # Check if the instantaneous wrench is correctly added to the permanent wrench cube_object.permanent_wrench_composer.add_forces_and_torques_index( @@ -304,8 +304,8 @@ def test_external_force_on_single_body(num_cubes, device): # Now we are ready! for i in range(5): # reset root state - root_pose = wp.to_torch(cube_object.data.default_root_pose).clone() - root_vel = wp.to_torch(cube_object.data.default_root_vel).clone() + root_pose = cube_object.data.default_root_pose.torch.clone() + root_vel = cube_object.data.default_root_vel.torch.clone() # need to shift the position of the cubes otherwise they will be on top of each other root_pose[:, :3] = origins @@ -318,7 +318,7 @@ def test_external_force_on_single_body(num_cubes, device): is_global = False if i % 2 == 0: is_global = True - positions = wp.to_torch(cube_object.data.body_com_pos_w)[:, body_ids, :3] + positions = cube_object.data.body_com_pos_w.torch[:, body_ids, :3] else: positions = None @@ -343,10 +343,10 @@ def test_external_force_on_single_body(num_cubes, device): # First object should still be at the same Z position (1.0) torch.testing.assert_close( - wp.to_torch(cube_object.data.root_pos_w)[0::2, 2], torch.ones(num_cubes // 2, device=sim.device) + cube_object.data.root_pos_w.torch[0::2, 2], torch.ones(num_cubes // 2, device=sim.device) ) # Second object should have fallen, so it's Z height should be less than initial height of 1.0 - assert torch.all(wp.to_torch(cube_object.data.root_pos_w)[1::2, 2] < 1.0) + assert torch.all(cube_object.data.root_pos_w.torch[1::2, 2] < 1.0) @pytest.mark.parametrize("num_cubes", [2, 4]) @@ -386,8 +386,8 @@ def test_external_force_on_single_body_at_position(num_cubes, device): # Now we are ready! for i in range(5): # reset root state - root_pose = wp.to_torch(cube_object.data.default_root_pose).clone() - root_vel = wp.to_torch(cube_object.data.default_root_vel).clone() + root_pose = cube_object.data.default_root_pose.torch.clone() + root_vel = cube_object.data.default_root_vel.torch.clone() # need to shift the position of the cubes otherwise they will be on top of each other root_pose[:, :3] = origins @@ -400,7 +400,7 @@ def test_external_force_on_single_body_at_position(num_cubes, device): is_global = False if i % 2 == 0: is_global = True - body_com_pos_w = wp.to_torch(cube_object.data.body_com_pos_w)[:, body_ids, :3] + body_com_pos_w = cube_object.data.body_com_pos_w.torch[:, body_ids, :3] external_wrench_positions_b[..., 0] = 0.0 external_wrench_positions_b[..., 1] = 1.0 external_wrench_positions_b[..., 2] = 0.0 @@ -426,13 +426,13 @@ def test_external_force_on_single_body_at_position(num_cubes, device): is_global=is_global, ) torch.testing.assert_close( - wp.to_torch(cube_object._permanent_wrench_composer.composed_force)[:, 0, :], + cube_object._permanent_wrench_composer.composed_force.torch[:, 0, :], desired_force[:, 0, :], rtol=1e-6, atol=1e-7, ) torch.testing.assert_close( - wp.to_torch(cube_object._permanent_wrench_composer.composed_torque)[:, 0, :], + cube_object._permanent_wrench_composer.composed_torque.torch[:, 0, :], desired_torque[:, 0, :], rtol=1e-6, atol=1e-7, @@ -449,9 +449,9 @@ def test_external_force_on_single_body_at_position(num_cubes, device): cube_object.update(sim.cfg.dt) # The first object should be rotating around it's X axis - assert torch.all(torch.abs(wp.to_torch(cube_object.data.root_ang_vel_b)[0::2, 0]) > 0.1) + assert torch.all(torch.abs(cube_object.data.root_ang_vel_b.torch[0::2, 0]) > 0.1) # Second object should have fallen, so it's Z height should be less than initial height of 1.0 - assert torch.all(wp.to_torch(cube_object.data.root_pos_w)[1::2, 2] < 1.0) + assert torch.all(cube_object.data.root_pos_w.torch[1::2, 2] < 1.0) @pytest.mark.parametrize("num_cubes", [1, 2]) @@ -479,10 +479,10 @@ def test_set_rigid_object_state(num_cubes, device): # Set each state type individually as they are dependent on each other for state_type_to_randomize in state_types: state_dict = { - "root_pos_w": torch.zeros_like(wp.to_torch(cube_object.data.root_pos_w), device=sim.device), + "root_pos_w": torch.zeros_like(cube_object.data.root_pos_w.torch, device=sim.device), "root_quat_w": default_orientation(num=num_cubes, device=sim.device), - "root_lin_vel_w": torch.zeros_like(wp.to_torch(cube_object.data.root_lin_vel_w), device=sim.device), - "root_ang_vel_w": torch.zeros_like(wp.to_torch(cube_object.data.root_ang_vel_w), device=sim.device), + "root_lin_vel_w": torch.zeros_like(cube_object.data.root_lin_vel_w.torch, device=sim.device), + "root_ang_vel_w": torch.zeros_like(cube_object.data.root_ang_vel_w.torch, device=sim.device), } # Now we are ready! @@ -514,7 +514,7 @@ def test_set_rigid_object_state(num_cubes, device): # assert that set root quantities are equal to the ones set in the state_dict for key, expected_value in state_dict.items(): - value = wp.to_torch(getattr(cube_object.data, key)) + value = getattr(cube_object.data, key).torch torch.testing.assert_close(value, expected_value, rtol=1e-3, atol=1e-3) cube_object.update(sim.cfg.dt) @@ -541,13 +541,13 @@ def test_reset_rigid_object(num_cubes, device): cube_object.update(sim.cfg.dt) # Move the object to a random position - root_pose = wp.to_torch(cube_object.data.default_root_pose).clone() + root_pose = cube_object.data.default_root_pose.torch.clone() root_pose[:, :3] = torch.randn(num_cubes, 3, device=sim.device) # Random orientation root_pose[:, 3:7] = random_orientation(num=num_cubes, device=sim.device) cube_object.write_root_pose_to_sim_index(root_pose=root_pose) - root_vel = wp.to_torch(cube_object.data.default_root_vel).clone() + root_vel = cube_object.data.default_root_vel.torch.clone() cube_object.write_root_velocity_to_sim_index(root_velocity=root_vel) if i % 2 == 0: @@ -557,10 +557,10 @@ def test_reset_rigid_object(num_cubes, device): # Reset should zero external forces and torques assert not cube_object._instantaneous_wrench_composer.active assert not cube_object._permanent_wrench_composer.active - assert torch.count_nonzero(wp.to_torch(cube_object._instantaneous_wrench_composer.composed_force)) == 0 - assert torch.count_nonzero(wp.to_torch(cube_object._instantaneous_wrench_composer.composed_torque)) == 0 - assert torch.count_nonzero(wp.to_torch(cube_object._permanent_wrench_composer.composed_force)) == 0 - assert torch.count_nonzero(wp.to_torch(cube_object._permanent_wrench_composer.composed_torque)) == 0 + assert torch.count_nonzero(cube_object._instantaneous_wrench_composer.composed_force.torch) == 0 + assert torch.count_nonzero(cube_object._instantaneous_wrench_composer.composed_torque.torch) == 0 + assert torch.count_nonzero(cube_object._permanent_wrench_composer.composed_force.torch) == 0 + assert torch.count_nonzero(cube_object._permanent_wrench_composer.composed_torque.torch) == 0 @pytest.mark.parametrize("num_cubes", [1, 2]) @@ -700,7 +700,7 @@ def test_rigid_body_no_friction(num_cubes, device): tolerance = 1e-5 torch.testing.assert_close( - wp.to_torch(cube_object.data.root_lin_vel_w), initial_velocity[:, :3], rtol=1e-5, atol=tolerance + cube_object.data.root_lin_vel_w.torch, initial_velocity[:, :3], rtol=1e-5, atol=tolerance ) @@ -773,7 +773,7 @@ def test_rigid_body_with_static_friction(num_cubes, device): ) # Get root state - initial_root_pos = wp.to_torch(cube_object.data.root_pos_w).clone() + initial_root_pos = cube_object.data.root_pos_w.torch.clone() # Simulate physics for _ in range(200): # apply the wrench @@ -784,10 +784,10 @@ def test_rigid_body_with_static_friction(num_cubes, device): if force == "below_mu": # Assert that the block has not moved torch.testing.assert_close( - wp.to_torch(cube_object.data.root_pos_w), initial_root_pos, rtol=2e-3, atol=2e-3 + cube_object.data.root_pos_w.torch, initial_root_pos, rtol=2e-3, atol=2e-3 ) if force == "above_mu": - assert (wp.to_torch(cube_object.data.root_pos_w)[..., 0] - initial_root_pos[..., 0] > 0.02).all() + assert (cube_object.data.root_pos_w.torch[..., 0] - initial_root_pos[..., 0] > 0.02).all() @pytest.mark.parametrize("num_cubes", [1, 2]) @@ -847,14 +847,14 @@ def test_rigid_body_with_restitution(num_cubes, device): wp.from_torch(cube_object_materials, dtype=wp.float32), wp.from_torch(indices, dtype=wp.int32) ) - curr_z_velocity = wp.to_torch(cube_object.data.root_lin_vel_w)[:, 2].clone() + curr_z_velocity = cube_object.data.root_lin_vel_w.torch[:, 2].clone() for _ in range(100): sim.step() # update object cube_object.update(sim.cfg.dt) - curr_z_velocity = wp.to_torch(cube_object.data.root_lin_vel_w)[:, 2].clone() + curr_z_velocity = cube_object.data.root_lin_vel_w.torch[:, 2].clone() if expected_collision_type == "inelastic": # assert that the block has not bounced by checking that the z velocity is less than or equal to 0 @@ -938,9 +938,9 @@ def test_gravity_vec_w(num_cubes, device, gravity_enabled): sim.reset() # Check that gravity is set correctly - assert wp.to_torch(cube_object.data.GRAVITY_VEC_W)[0, 0] == gravity_dir[0] - assert wp.to_torch(cube_object.data.GRAVITY_VEC_W)[0, 1] == gravity_dir[1] - assert wp.to_torch(cube_object.data.GRAVITY_VEC_W)[0, 2] == gravity_dir[2] + assert cube_object.data.GRAVITY_VEC_W.torch[0, 0] == gravity_dir[0] + assert cube_object.data.GRAVITY_VEC_W.torch[0, 1] == gravity_dir[1] + assert cube_object.data.GRAVITY_VEC_W.torch[0, 2] == gravity_dir[2] # Simulate physics for _ in range(2): @@ -954,7 +954,7 @@ def test_gravity_vec_w(num_cubes, device, gravity_enabled): if gravity_enabled: gravity[:, :, 2] = -9.81 # Check the body accelerations are correct - torch.testing.assert_close(wp.to_torch(cube_object.data.body_acc_w), gravity) + torch.testing.assert_close(cube_object.data.body_acc_w.torch, gravity) @pytest.mark.parametrize("num_cubes", [1, 2]) @@ -1003,14 +1003,14 @@ def test_body_root_state_properties(num_cubes, device, with_offset): cube_object.update(sim.cfg.dt) # get state properties - root_link_pose_w = wp.to_torch(cube_object.data.root_link_pose_w) - root_link_vel_w = wp.to_torch(cube_object.data.root_link_vel_w) - root_com_pose_w = wp.to_torch(cube_object.data.root_com_pose_w) - root_com_vel_w = wp.to_torch(cube_object.data.root_com_vel_w) - body_link_pose_w = wp.to_torch(cube_object.data.body_link_pose_w) - body_link_vel_w = wp.to_torch(cube_object.data.body_link_vel_w) - body_com_pose_w = wp.to_torch(cube_object.data.body_com_pose_w) - body_com_vel_w = wp.to_torch(cube_object.data.body_com_vel_w) + root_link_pose_w = cube_object.data.root_link_pose_w.torch + root_link_vel_w = cube_object.data.root_link_vel_w.torch + root_com_pose_w = cube_object.data.root_com_pose_w.torch + root_com_vel_w = cube_object.data.root_com_vel_w.torch + body_link_pose_w = cube_object.data.body_link_pose_w.torch + body_link_vel_w = cube_object.data.body_link_vel_w.torch + body_com_pose_w = cube_object.data.body_com_pose_w.torch + body_com_vel_w = cube_object.data.body_com_vel_w.torch # if offset is [0,0,0] all root_state_%_w will match and all body_%_w will match if not with_offset: @@ -1041,7 +1041,7 @@ def test_body_root_state_properties(num_cubes, device, with_offset): torch.testing.assert_close(-offset, body_link_state_pos_rel_com.squeeze(-2)) # orientation of com will be a constant rotation from link orientation - com_quat_b = wp.to_torch(cube_object.data.body_com_quat_b) + com_quat_b = cube_object.data.body_com_quat_b.torch com_quat_w = quat_mul(body_link_pose_w[..., 3:], com_quat_b) torch.testing.assert_close(com_quat_w, body_com_pose_w[..., 3:]) torch.testing.assert_close(com_quat_w.squeeze(-2), root_com_pose_w[..., 3:]) @@ -1101,7 +1101,7 @@ def test_write_root_state(num_cubes, device, with_offset, state_location): torch.testing.assert_close(wp.to_torch(cube_object.root_view.get_coms()), com) rand_state = torch.zeros(num_cubes, 13, device=device) - rand_state[..., :7] = wp.to_torch(cube_object.data.default_root_pose) + rand_state[..., :7] = cube_object.data.default_root_pose.torch rand_state[..., :3] += env_pos # make quaternion a unit vector rand_state[..., 3:7] = torch.nn.functional.normalize(rand_state[..., 3:7], dim=-1) @@ -1131,11 +1131,11 @@ def test_write_root_state(num_cubes, device, with_offset, state_location): ) if state_location == "com": - torch.testing.assert_close(rand_state[..., :7], wp.to_torch(cube_object.data.root_com_pose_w)) - torch.testing.assert_close(rand_state[..., 7:], wp.to_torch(cube_object.data.root_com_vel_w)) + torch.testing.assert_close(rand_state[..., :7], cube_object.data.root_com_pose_w.torch) + torch.testing.assert_close(rand_state[..., 7:], cube_object.data.root_com_vel_w.torch) elif state_location == "link": - torch.testing.assert_close(rand_state[..., :7], wp.to_torch(cube_object.data.root_link_pose_w)) - torch.testing.assert_close(rand_state[..., 7:], wp.to_torch(cube_object.data.root_link_vel_w)) + torch.testing.assert_close(rand_state[..., :7], cube_object.data.root_link_pose_w.torch) + torch.testing.assert_close(rand_state[..., 7:], cube_object.data.root_link_vel_w.torch) @pytest.mark.parametrize("num_cubes", [1, 2]) @@ -1193,9 +1193,9 @@ def test_write_state_functions_data_consistency(num_cubes, device, with_offset, cube_object.write_root_velocity_to_sim_index(root_velocity=rand_state[..., 7:]) if state_location == "com": - root_com_pose_w = wp.to_torch(cube_object.data.root_com_pose_w) - root_com_vel_w = wp.to_torch(cube_object.data.root_com_vel_w) - body_com_pose_b = wp.to_torch(cube_object.data.body_com_pose_b) + root_com_pose_w = cube_object.data.root_com_pose_w.torch + root_com_vel_w = cube_object.data.root_com_vel_w.torch + body_com_pose_b = cube_object.data.body_com_pose_b.torch expected_root_link_pos, expected_root_link_quat = combine_frame_transforms( root_com_pose_w[:, :3], root_com_pose_w[:, 3:], @@ -1203,19 +1203,19 @@ def test_write_state_functions_data_consistency(num_cubes, device, with_offset, quat_inv(body_com_pose_b[:, 0, 3:7]), ) expected_root_link_pose = torch.cat((expected_root_link_pos, expected_root_link_quat), dim=1) - root_link_pose_w = wp.to_torch(cube_object.data.root_link_pose_w) - root_link_vel_w = wp.to_torch(cube_object.data.root_link_vel_w) + root_link_pose_w = cube_object.data.root_link_pose_w.torch + root_link_vel_w = cube_object.data.root_link_vel_w.torch # test both root_pose and root_link successfully updated when root_com updates torch.testing.assert_close(expected_root_link_pose, root_link_pose_w) # skip lin_vel because it differs from link frame, this should be fine because we are only checking # if velocity update is triggered, which can be determined by comparing angular velocity torch.testing.assert_close(root_com_vel_w[:, 3:], root_link_vel_w[:, 3:]) torch.testing.assert_close(expected_root_link_pose, root_link_pose_w) - torch.testing.assert_close(root_com_vel_w[:, 3:], wp.to_torch(cube_object.data.root_com_vel_w)[:, 3:]) + torch.testing.assert_close(root_com_vel_w[:, 3:], cube_object.data.root_com_vel_w.torch[:, 3:]) elif state_location == "link": - root_link_pose_w = wp.to_torch(cube_object.data.root_link_pose_w) - root_link_vel_w = wp.to_torch(cube_object.data.root_link_vel_w) - body_com_pose_b = wp.to_torch(cube_object.data.body_com_pose_b) + root_link_pose_w = cube_object.data.root_link_pose_w.torch + root_link_vel_w = cube_object.data.root_link_vel_w.torch + body_com_pose_b = cube_object.data.body_com_pose_b.torch expected_com_pos, expected_com_quat = combine_frame_transforms( root_link_pose_w[:, :3], root_link_pose_w[:, 3:], @@ -1223,19 +1223,19 @@ def test_write_state_functions_data_consistency(num_cubes, device, with_offset, body_com_pose_b[:, 0, 3:7], ) expected_com_pose = torch.cat((expected_com_pos, expected_com_quat), dim=1) - root_com_pose_w = wp.to_torch(cube_object.data.root_com_pose_w) - root_com_vel_w = wp.to_torch(cube_object.data.root_com_vel_w) + root_com_pose_w = cube_object.data.root_com_pose_w.torch + root_com_vel_w = cube_object.data.root_com_vel_w.torch # test both root_pose and root_com successfully updated when root_link updates torch.testing.assert_close(expected_com_pose, root_com_pose_w) # skip lin_vel because it differs from link frame, this should be fine because we are only checking # if velocity update is triggered, which can be determined by comparing angular velocity torch.testing.assert_close(root_link_vel_w[:, 3:], root_com_vel_w[:, 3:]) - torch.testing.assert_close(root_link_pose_w, wp.to_torch(cube_object.data.root_link_pose_w)) - torch.testing.assert_close(root_link_vel_w[:, 3:], wp.to_torch(cube_object.data.root_com_vel_w)[:, 3:]) + torch.testing.assert_close(root_link_pose_w, cube_object.data.root_link_pose_w.torch) + torch.testing.assert_close(root_link_vel_w[:, 3:], cube_object.data.root_com_vel_w.torch[:, 3:]) elif state_location == "root": - root_link_pose_w = wp.to_torch(cube_object.data.root_link_pose_w) - root_com_vel_w = wp.to_torch(cube_object.data.root_com_vel_w) - body_com_pose_b = wp.to_torch(cube_object.data.body_com_pose_b) + root_link_pose_w = cube_object.data.root_link_pose_w.torch + root_com_vel_w = cube_object.data.root_com_vel_w.torch + body_com_pose_b = cube_object.data.body_com_pose_b.torch expected_com_pos, expected_com_quat = combine_frame_transforms( root_link_pose_w[:, :3], root_link_pose_w[:, 3:], @@ -1243,12 +1243,12 @@ def test_write_state_functions_data_consistency(num_cubes, device, with_offset, body_com_pose_b[:, 0, 3:7], ) expected_com_pose = torch.cat((expected_com_pos, expected_com_quat), dim=1) - root_com_pose_w = wp.to_torch(cube_object.data.root_com_pose_w) - root_link_vel_w = wp.to_torch(cube_object.data.root_link_vel_w) + root_com_pose_w = cube_object.data.root_com_pose_w.torch + root_link_vel_w = cube_object.data.root_link_vel_w.torch # test both root_com and root_link successfully updated when root_pose updates torch.testing.assert_close(expected_com_pose, root_com_pose_w) - torch.testing.assert_close(root_com_vel_w, wp.to_torch(cube_object.data.root_com_vel_w)) - torch.testing.assert_close(root_link_pose_w, wp.to_torch(cube_object.data.root_link_pose_w)) + torch.testing.assert_close(root_com_vel_w, cube_object.data.root_com_vel_w.torch) + torch.testing.assert_close(root_link_pose_w, cube_object.data.root_link_pose_w.torch) torch.testing.assert_close(root_com_vel_w[:, 3:], root_link_vel_w[:, 3:]) diff --git a/source/isaaclab_physx/test/assets/test_rigid_object_collection.py b/source/isaaclab_physx/test/assets/test_rigid_object_collection.py index 03689d73f48a..8401cc395a0c 100644 --- a/source/isaaclab_physx/test/assets/test_rigid_object_collection.py +++ b/source/isaaclab_physx/test/assets/test_rigid_object_collection.py @@ -126,10 +126,10 @@ def test_initialization(sim, num_envs, num_cubes, device): assert len(object_collection.body_names) == num_cubes # Check buffers that exist and have correct shapes - assert wp.to_torch(object_collection.data.body_link_pos_w).shape == (num_envs, num_cubes, 3) - assert wp.to_torch(object_collection.data.body_link_quat_w).shape == (num_envs, num_cubes, 4) - assert wp.to_torch(object_collection.data.body_mass).shape == (num_envs, num_cubes) - assert wp.to_torch(object_collection.data.body_inertia).shape == (num_envs, num_cubes, 9) + assert object_collection.data.body_link_pos_w.torch.shape == (num_envs, num_cubes, 3) + assert object_collection.data.body_link_quat_w.torch.shape == (num_envs, num_cubes, 4) + assert object_collection.data.body_mass.torch.shape == (num_envs, num_cubes) + assert object_collection.data.body_inertia.torch.shape == (num_envs, num_cubes, 9) # Simulate physics for _ in range(2): @@ -193,19 +193,19 @@ def test_initialization_with_kinematic_enabled(sim, num_envs, num_cubes, device) assert len(object_collection.body_names) == num_cubes # Check buffers that exist and have correct shapes - assert wp.to_torch(object_collection.data.body_link_pos_w).shape == (num_envs, num_cubes, 3) - assert wp.to_torch(object_collection.data.body_link_quat_w).shape == (num_envs, num_cubes, 4) + assert object_collection.data.body_link_pos_w.torch.shape == (num_envs, num_cubes, 3) + assert object_collection.data.body_link_quat_w.torch.shape == (num_envs, num_cubes, 4) # Simulate physics for _ in range(2): sim.step() object_collection.update(sim.cfg.dt) # check that the object is kinematic - default_body_pose = wp.to_torch(object_collection.data.default_body_pose).clone() - default_body_vel = wp.to_torch(object_collection.data.default_body_vel).clone() + default_body_pose = object_collection.data.default_body_pose.torch.clone() + default_body_vel = object_collection.data.default_body_vel.torch.clone() default_body_pose[..., :3] += origins.unsqueeze(1) - torch.testing.assert_close(wp.to_torch(object_collection.data.body_link_pose_w), default_body_pose) - torch.testing.assert_close(wp.to_torch(object_collection.data.body_link_vel_w), default_body_vel) + torch.testing.assert_close(object_collection.data.body_link_pose_w.torch, default_body_pose) + torch.testing.assert_close(object_collection.data.body_link_vel_w.torch, default_body_vel) @pytest.mark.parametrize("num_cubes", [1, 2]) @@ -259,8 +259,8 @@ def test_external_force_buffer(sim, device): # check if the object collection's force and torque buffers are correctly updated for i in range(num_envs): - assert wp.to_torch(object_collection._permanent_wrench_composer.composed_force)[i, 0, 0].item() == force - assert wp.to_torch(object_collection._permanent_wrench_composer.composed_torque)[i, 0, 0].item() == force + assert object_collection._permanent_wrench_composer.composed_force.torch[i, 0, 0].item() == force + assert object_collection._permanent_wrench_composer.composed_torque.torch[i, 0, 0].item() == force object_collection.instantaneous_wrench_composer.add_forces_and_torques_index( body_ids=object_ids, @@ -288,12 +288,12 @@ def test_external_force_on_single_body(sim, num_envs, num_cubes, device): # Sample a force equal to the weight of the object external_wrench_b = torch.zeros(object_collection.num_instances, len(object_ids), 6, device=sim.device) # Every 2nd cube should have a force applied to it - external_wrench_b[:, 0::2, 2] = 9.81 * wp.to_torch(object_collection.data.body_mass)[:, 0::2] + external_wrench_b[:, 0::2, 2] = 9.81 * object_collection.data.body_mass.torch[:, 0::2] for i in range(5): # reset object state - body_pose = wp.to_torch(object_collection.data.default_body_pose).clone() - body_vel = wp.to_torch(object_collection.data.default_body_vel).clone() + body_pose = object_collection.data.default_body_pose.torch.clone() + body_vel = object_collection.data.default_body_vel.torch.clone() # need to shift the position of the cubes otherwise they will be on top of each other body_pose[..., :2] += origins.unsqueeze(1)[..., :2] object_collection.write_body_link_pose_to_sim_index(body_poses=body_pose) @@ -303,7 +303,7 @@ def test_external_force_on_single_body(sim, num_envs, num_cubes, device): is_global = False if i % 2 == 0: - positions = wp.to_torch(object_collection.data.body_link_pos_w)[:, object_ids, :3] + positions = object_collection.data.body_link_pos_w.torch[:, object_ids, :3] is_global = True else: positions = None @@ -327,11 +327,11 @@ def test_external_force_on_single_body(sim, num_envs, num_cubes, device): # First object should still be at the same Z position (1.0) torch.testing.assert_close( - wp.to_torch(object_collection.data.body_link_pos_w)[:, 0::2, 2], - torch.ones_like(wp.to_torch(object_collection.data.body_link_pos_w)[:, 0::2, 2]), + object_collection.data.body_link_pos_w.torch[:, 0::2, 2], + torch.ones_like(object_collection.data.body_link_pos_w.torch[:, 0::2, 2]), ) # Second object should have fallen, so it's Z height should be less than initial height of 1.0 - assert torch.all(wp.to_torch(object_collection.data.body_link_pos_w)[:, 1::2, 2] < 1.0) + assert torch.all(object_collection.data.body_link_pos_w.torch[:, 1::2, 2] < 1.0) @pytest.mark.parametrize("num_envs", [1, 2]) @@ -360,8 +360,8 @@ def test_external_force_on_single_body_at_position(sim, num_envs, num_cubes, dev # Desired force and torque for i in range(5): # reset object state - body_pose = wp.to_torch(object_collection.data.default_body_pose).clone() - body_vel = wp.to_torch(object_collection.data.default_body_vel).clone() + body_pose = object_collection.data.default_body_pose.torch.clone() + body_vel = object_collection.data.default_body_vel.torch.clone() # need to shift the position of the cubes otherwise they will be on top of each other body_pose[..., :2] += origins.unsqueeze(1)[..., :2] object_collection.write_body_link_pose_to_sim_index(body_poses=body_pose) @@ -371,7 +371,7 @@ def test_external_force_on_single_body_at_position(sim, num_envs, num_cubes, dev is_global = False if i % 2 == 0: - body_com_pos_w = wp.to_torch(object_collection.data.body_link_pos_w)[:, object_ids, :3] + body_com_pos_w = object_collection.data.body_link_pos_w.torch[:, object_ids, :3] external_wrench_positions_b[..., 0] = 0.0 external_wrench_positions_b[..., 1] = 1.0 external_wrench_positions_b[..., 2] = 0.0 @@ -408,9 +408,9 @@ def test_external_force_on_single_body_at_position(sim, num_envs, num_cubes, dev object_collection.update(sim.cfg.dt) # First object should be rotating around it's X axis - assert torch.all(wp.to_torch(object_collection.data.body_com_ang_vel_b)[:, 0::2, 0] > 0.1) + assert torch.all(object_collection.data.body_com_ang_vel_b.torch[:, 0::2, 0] > 0.1) # Second object should have fallen, so it's Z height should be less than initial height of 1.0 - assert torch.all(wp.to_torch(object_collection.data.body_link_pos_w)[:, 1::2, 2] < 1.0) + assert torch.all(object_collection.data.body_link_pos_w.torch[:, 1::2, 2] < 1.0) @pytest.mark.parametrize("num_envs", [1, 3]) @@ -432,16 +432,12 @@ def test_set_object_state(sim, num_envs, num_cubes, device, gravity_enabled): # Set each state type individually as they are dependent on each other for state_type_to_randomize in state_types: state_dict = { - "body_link_pos_w": torch.zeros_like(wp.to_torch(object_collection.data.body_link_pos_w), device=sim.device), + "body_link_pos_w": torch.zeros_like(object_collection.data.body_link_pos_w.torch, device=sim.device), "body_link_quat_w": default_orientation(num=num_cubes * num_envs, device=sim.device).view( num_envs, num_cubes, 4 ), - "body_com_lin_vel_w": torch.zeros_like( - wp.to_torch(object_collection.data.body_com_lin_vel_w), device=sim.device - ), - "body_com_ang_vel_w": torch.zeros_like( - wp.to_torch(object_collection.data.body_com_ang_vel_w), device=sim.device - ), + "body_com_lin_vel_w": torch.zeros_like(object_collection.data.body_com_lin_vel_w.torch, device=sim.device), + "body_com_ang_vel_w": torch.zeros_like(object_collection.data.body_com_ang_vel_w.torch, device=sim.device), } for _ in range(5): @@ -476,7 +472,7 @@ def test_set_object_state(sim, num_envs, num_cubes, device, gravity_enabled): # assert that set object quantities are equal to the ones set in the state_dict for key, expected_value in state_dict.items(): - value = wp.to_torch(getattr(object_collection.data, key)) + value = getattr(object_collection.data, key).torch torch.testing.assert_close(value, expected_value, rtol=1e-5, atol=1e-5) object_collection.update(sim.cfg.dt) @@ -521,7 +517,7 @@ def test_object_state_properties(sim, num_envs, num_cubes, device, with_offset, spin_twist[5] = torch.randn(1, device=device) # initial spawn point - init_com = wp.to_torch(cube_object.data.body_com_pose_w)[..., :3] + init_com = cube_object.data.body_com_pose_w.torch[..., :3] for i in range(10): # spin the object around Z axis (com) @@ -530,10 +526,10 @@ def test_object_state_properties(sim, num_envs, num_cubes, device, with_offset, cube_object.update(sim.cfg.dt) # get state properties - object_link_pose_w = wp.to_torch(cube_object.data.body_link_pose_w) - object_link_vel_w = wp.to_torch(cube_object.data.body_link_vel_w) - object_com_pose_w = wp.to_torch(cube_object.data.body_com_pose_w) - object_com_vel_w = wp.to_torch(cube_object.data.body_com_vel_w) + object_link_pose_w = cube_object.data.body_link_pose_w.torch + object_link_vel_w = cube_object.data.body_link_vel_w.torch + object_com_pose_w = cube_object.data.body_com_pose_w.torch + object_com_vel_w = cube_object.data.body_com_vel_w.torch # if offset is [0,0,0] all object_state_%_w will match and all body_%_w will match if not with_offset: @@ -554,7 +550,7 @@ def test_object_state_properties(sim, num_envs, num_cubes, device, with_offset, torch.testing.assert_close(-offset, object_link_state_pos_rel_com) # orientation of com will be a constant rotation from link orientation - com_quat_b = wp.to_torch(cube_object.data.body_com_quat_b) + com_quat_b = cube_object.data.body_com_quat_b.torch com_quat_w = quat_mul(object_link_pose_w[..., 3:], com_quat_b) torch.testing.assert_close(com_quat_w, object_com_pose_w[..., 3:]) @@ -616,8 +612,8 @@ def test_write_object_state(sim, num_envs, num_cubes, device, with_offset, state ) rand_state = torch.zeros(num_envs, num_cubes, 13, device=device) - rand_state[..., :7] = wp.to_torch(cube_object.data.default_body_pose) - rand_state[..., :3] += wp.to_torch(cube_object.data.body_link_pos_w) + rand_state[..., :7] = cube_object.data.default_body_pose.torch + rand_state[..., :3] += cube_object.data.body_link_pos_w.torch # make quaternion a unit vector rand_state[..., 3:7] = torch.nn.functional.normalize(rand_state[..., 3:7], dim=-1) @@ -651,11 +647,11 @@ def test_write_object_state(sim, num_envs, num_cubes, device, with_offset, state ) if state_location == "com": - torch.testing.assert_close(rand_state[..., :7], wp.to_torch(cube_object.data.body_com_pose_w)) - torch.testing.assert_close(rand_state[..., 7:], wp.to_torch(cube_object.data.body_com_vel_w)) + torch.testing.assert_close(rand_state[..., :7], cube_object.data.body_com_pose_w.torch) + torch.testing.assert_close(rand_state[..., 7:], cube_object.data.body_com_vel_w.torch) elif state_location == "link": - torch.testing.assert_close(rand_state[..., :7], wp.to_torch(cube_object.data.body_link_pose_w)) - torch.testing.assert_close(rand_state[..., 7:], wp.to_torch(cube_object.data.body_link_vel_w)) + torch.testing.assert_close(rand_state[..., :7], cube_object.data.body_link_pose_w.torch) + torch.testing.assert_close(rand_state[..., 7:], cube_object.data.body_link_vel_w.torch) @pytest.mark.parametrize("num_envs", [1, 3]) @@ -671,12 +667,12 @@ def test_reset_object_collection(sim, num_envs, num_cubes, device): object_collection.update(sim.cfg.dt) # Move the object to a random position - body_pose = wp.to_torch(object_collection.data.default_body_pose).clone() + body_pose = object_collection.data.default_body_pose.torch.clone() body_pose[..., :3] = torch.randn(num_envs, num_cubes, 3, device=sim.device) # Random orientation body_pose[..., 3:7] = random_orientation(num=num_cubes, device=sim.device) object_collection.write_body_link_pose_to_sim_index(body_poses=body_pose) - body_vel = wp.to_torch(object_collection.data.default_body_vel).clone() + body_vel = object_collection.data.default_body_vel.torch.clone() object_collection.write_body_com_velocity_to_sim_index(body_velocities=body_vel) if i % 2 == 0: @@ -685,14 +681,10 @@ def test_reset_object_collection(sim, num_envs, num_cubes, device): # Reset should zero external forces and torques assert not object_collection._instantaneous_wrench_composer.active assert not object_collection._permanent_wrench_composer.active - assert ( - torch.count_nonzero(wp.to_torch(object_collection._instantaneous_wrench_composer.composed_force)) == 0 - ) - assert ( - torch.count_nonzero(wp.to_torch(object_collection._instantaneous_wrench_composer.composed_torque)) == 0 - ) - assert torch.count_nonzero(wp.to_torch(object_collection._permanent_wrench_composer.composed_force)) == 0 - assert torch.count_nonzero(wp.to_torch(object_collection._permanent_wrench_composer.composed_torque)) == 0 + assert torch.count_nonzero(object_collection._instantaneous_wrench_composer.composed_force.torch) == 0 + assert torch.count_nonzero(object_collection._instantaneous_wrench_composer.composed_torque.torch) == 0 + assert torch.count_nonzero(object_collection._permanent_wrench_composer.composed_force.torch) == 0 + assert torch.count_nonzero(object_collection._permanent_wrench_composer.composed_torque.torch) == 0 @pytest.mark.parametrize("num_envs", [1, 3]) @@ -744,7 +736,7 @@ def test_gravity_vec_w(sim, num_envs, num_cubes, device, gravity_enabled): sim.reset() # Check if gravity vector is set correctly - gravity_vec = wp.to_torch(object_collection.data.GRAVITY_VEC_W) + gravity_vec = object_collection.data.GRAVITY_VEC_W.torch assert gravity_vec[0, 0, 0] == gravity_dir[0] assert gravity_vec[0, 0, 1] == gravity_dir[1] assert gravity_vec[0, 0, 2] == gravity_dir[2] @@ -760,7 +752,7 @@ def test_gravity_vec_w(sim, num_envs, num_cubes, device, gravity_enabled): gravity[..., 2] = -9.81 # Check the body accelerations are correct - torch.testing.assert_close(wp.to_torch(object_collection.data.body_com_acc_w), gravity) + torch.testing.assert_close(object_collection.data.body_com_acc_w.torch, gravity) @pytest.mark.parametrize("num_envs", [1, 3]) @@ -804,7 +796,7 @@ def test_write_object_state_functions_data_consistency( ) rand_state = torch.rand(num_envs, num_cubes, 13, device=device) - rand_state[..., :3] += wp.to_torch(cube_object.data.body_link_pos_w) + rand_state[..., :3] += cube_object.data.body_link_pos_w.torch # make quaternion a unit vector rand_state[..., 3:7] = torch.nn.functional.normalize(rand_state[..., 3:7], dim=-1) @@ -813,8 +805,8 @@ def test_write_object_state_functions_data_consistency( sim.step() cube_object.update(sim.cfg.dt) - body_link_pose_w = wp.to_torch(cube_object.data.body_link_pose_w) - body_com_pose_w = wp.to_torch(cube_object.data.body_com_pose_w) + body_link_pose_w = cube_object.data.body_link_pose_w.torch + body_com_pose_w = cube_object.data.body_com_pose_w.torch object_link_to_com_pos, object_link_to_com_quat = subtract_frame_transforms( body_link_pose_w[..., :3].view(-1, 3), body_link_pose_w[..., 3:7].view(-1, 4), @@ -845,8 +837,8 @@ def test_write_object_state_functions_data_consistency( ) if state_location == "com": - com_pose_w = wp.to_torch(cube_object.data.body_com_pose_w) - com_vel_w = wp.to_torch(cube_object.data.body_com_vel_w) + com_pose_w = cube_object.data.body_com_pose_w.torch + com_vel_w = cube_object.data.body_com_vel_w.torch expected_root_link_pos, expected_root_link_quat = combine_frame_transforms( com_pose_w[..., :3].view(-1, 3), com_pose_w[..., 3:].view(-1, 4), @@ -856,18 +848,18 @@ def test_write_object_state_functions_data_consistency( expected_object_link_pose = torch.cat((expected_root_link_pos, expected_root_link_quat), dim=1).view( num_envs, -1, 7 ) - link_pose_w = wp.to_torch(cube_object.data.body_link_pose_w) - link_vel_w = wp.to_torch(cube_object.data.body_link_vel_w) + link_pose_w = cube_object.data.body_link_pose_w.torch + link_vel_w = cube_object.data.body_link_vel_w.torch # test both root_pose and root_link successfully updated when root_com updates torch.testing.assert_close(expected_object_link_pose, link_pose_w) # skip lin_vel because it differs from link frame, this should be fine because we are only checking # if velocity update is triggered, which can be determined by comparing angular velocity torch.testing.assert_close(com_vel_w[..., 3:], link_vel_w[..., 3:]) torch.testing.assert_close(expected_object_link_pose, link_pose_w) - torch.testing.assert_close(com_vel_w[..., 3:], wp.to_torch(cube_object.data.body_com_vel_w)[..., 3:]) + torch.testing.assert_close(com_vel_w[..., 3:], cube_object.data.body_com_vel_w.torch[..., 3:]) elif state_location == "link": - link_pose_w = wp.to_torch(cube_object.data.body_link_pose_w) - link_vel_w = wp.to_torch(cube_object.data.body_link_vel_w) + link_pose_w = cube_object.data.body_link_pose_w.torch + link_vel_w = cube_object.data.body_link_vel_w.torch expected_com_pos, expected_com_quat = combine_frame_transforms( link_pose_w[..., :3].view(-1, 3), link_pose_w[..., 3:].view(-1, 4), @@ -875,18 +867,18 @@ def test_write_object_state_functions_data_consistency( object_link_to_com_quat, ) expected_object_com_pose = torch.cat((expected_com_pos, expected_com_quat), dim=1).view(num_envs, -1, 7) - com_pose_w = wp.to_torch(cube_object.data.body_com_pose_w) - com_vel_w = wp.to_torch(cube_object.data.body_com_vel_w) + com_pose_w = cube_object.data.body_com_pose_w.torch + com_vel_w = cube_object.data.body_com_vel_w.torch # test both root_pose and root_com successfully updated when root_link updates torch.testing.assert_close(expected_object_com_pose, com_pose_w) # skip lin_vel because it differs from link frame, this should be fine because we are only checking # if velocity update is triggered, which can be determined by comparing angular velocity torch.testing.assert_close(link_vel_w[..., 3:], com_vel_w[..., 3:]) - torch.testing.assert_close(link_pose_w, wp.to_torch(cube_object.data.body_link_pose_w)) - torch.testing.assert_close(link_vel_w[..., 3:], wp.to_torch(cube_object.data.body_com_vel_w)[..., 3:]) + torch.testing.assert_close(link_pose_w, cube_object.data.body_link_pose_w.torch) + torch.testing.assert_close(link_vel_w[..., 3:], cube_object.data.body_com_vel_w.torch[..., 3:]) elif state_location == "root": - body_link_pose_w = wp.to_torch(cube_object.data.body_link_pose_w) - body_com_vel_w = wp.to_torch(cube_object.data.body_com_vel_w) + body_link_pose_w = cube_object.data.body_link_pose_w.torch + body_com_vel_w = cube_object.data.body_com_vel_w.torch expected_object_com_pos, expected_object_com_quat = combine_frame_transforms( body_link_pose_w[..., :3].view(-1, 3), body_link_pose_w[..., 3:].view(-1, 4), @@ -896,10 +888,10 @@ def test_write_object_state_functions_data_consistency( expected_object_com_pose = torch.cat((expected_object_com_pos, expected_object_com_quat), dim=1).view( num_envs, -1, 7 ) - com_pose_w = wp.to_torch(cube_object.data.body_com_pose_w) - com_vel_w = wp.to_torch(cube_object.data.body_com_vel_w) - link_pose_w = wp.to_torch(cube_object.data.body_link_pose_w) - link_vel_w = wp.to_torch(cube_object.data.body_link_vel_w) + com_pose_w = cube_object.data.body_com_pose_w.torch + com_vel_w = cube_object.data.body_com_vel_w.torch + link_pose_w = cube_object.data.body_link_pose_w.torch + link_vel_w = cube_object.data.body_link_vel_w.torch # test both root_com and root_link successfully updated when root_pose updates torch.testing.assert_close(expected_object_com_pose, com_pose_w) torch.testing.assert_close(body_com_vel_w, com_vel_w) diff --git a/source/isaaclab_physx/test/sensors/check_pva_sensor.py b/source/isaaclab_physx/test/sensors/check_pva_sensor.py index 75bc64673df3..41476babf542 100644 --- a/source/isaaclab_physx/test/sensors/check_pva_sensor.py +++ b/source/isaaclab_physx/test/sensors/check_pva_sensor.py @@ -147,8 +147,8 @@ def main(): # Get the ball initial positions sim.step(render=not args_cli.headless) balls.update(sim.get_physics_dt()) - ball_initial_positions = balls.data.root_pos_w.clone() - ball_initial_orientations = balls.data.root_quat_w.clone() + ball_initial_positions = balls.data.root_pos_w.torch.clone() + ball_initial_orientations = balls.data.root_quat_w.torch.clone() # Create a counter for resetting the scene step_count = 0 diff --git a/source/isaaclab_physx/test/sensors/test_contact_sensor.py b/source/isaaclab_physx/test/sensors/test_contact_sensor.py index 0d7ea255006e..00772e6cb0d3 100644 --- a/source/isaaclab_physx/test/sensors/test_contact_sensor.py +++ b/source/isaaclab_physx/test/sensors/test_contact_sensor.py @@ -300,17 +300,17 @@ def test_cube_stack_contact_filtering(setup_simulation, device, num_envs): # Check values for cube 2 --> cube 1 is the only collision for cube 2 torch.testing.assert_close( - wp.to_torch(contact_sensor_2.data.force_matrix_w)[:, :, 0], - wp.to_torch(contact_sensor_2.data.net_forces_w), + contact_sensor_2.data.force_matrix_w.torch[:, :, 0], + contact_sensor_2.data.net_forces_w.torch, ) # Check that forces are opposite and equal torch.testing.assert_close( - wp.to_torch(contact_sensor_2.data.force_matrix_w)[:, :, 0], - -wp.to_torch(contact_sensor.data.force_matrix_w)[:, :, 0], + contact_sensor_2.data.force_matrix_w.torch[:, :, 0], + -contact_sensor.data.force_matrix_w.torch[:, :, 0], ) # Check values are non-zero (contacts are happening and are getting reported) - assert wp.to_torch(contact_sensor_2.data.net_forces_w).sum().item() > 0.0 - assert wp.to_torch(contact_sensor.data.net_forces_w).sum().item() > 0.0 + assert contact_sensor_2.data.net_forces_w.torch.sum().item() > 0.0 + assert contact_sensor.data.net_forces_w.torch.sum().item() > 0.0 def test_no_contact_reporting(setup_simulation): @@ -370,10 +370,10 @@ def test_no_contact_reporting(setup_simulation): _perform_sim_step(sim, scene, sim_dt) # check values are zero (contacts are happening but not reported) - assert wp.to_torch(contact_sensor.data.net_forces_w).sum().item() == 0.0 - assert wp.to_torch(contact_sensor.data.force_matrix_w).sum().item() == 0.0 - assert wp.to_torch(contact_sensor_2.data.net_forces_w).sum().item() == 0.0 - assert wp.to_torch(contact_sensor_2.data.force_matrix_w).sum().item() == 0.0 + assert contact_sensor.data.net_forces_w.torch.sum().item() == 0.0 + assert contact_sensor.data.force_matrix_w.torch.sum().item() == 0.0 + assert contact_sensor_2.data.net_forces_w.torch.sum().item() == 0.0 + assert contact_sensor_2.data.force_matrix_w.torch.sum().item() == 0.0 @pytest.mark.isaacsim_ci @@ -490,7 +490,7 @@ def test_friction_reporting(setup_simulation, grav_dir): # check that forces are being reported match expected friction forces expected_friction, _, _, _ = scene["contact_sensor"].contact_view.get_friction_data(dt=sim_dt) expected_friction_torch = wp.to_torch(expected_friction) - reported_friction = wp.to_torch(scene["contact_sensor"].data.friction_forces_w)[0, 0, :] + reported_friction = scene["contact_sensor"].data.friction_forces_w.torch[0, 0, :] torch.testing.assert_close(expected_friction_torch.sum(dim=0), reported_friction[0], atol=1e-6, rtol=1e-5) @@ -751,7 +751,7 @@ def _test_friction_forces(shape: RigidObject, sensor: ContactSensor, mode: Conta # check shape of the friction_forces_w tensor (wp.to_torch expands vec3f -> float32 trailing dim) num_bodies = sensor.num_bodies - friction_torch = wp.to_torch(sensor._data.friction_forces_w) + friction_torch = sensor._data.friction_forces_w.torch assert friction_torch.shape == (sensor.num_instances // num_bodies, num_bodies, 1, 3) # compare friction forces if mode == ContactTestMode.IN_CONTACT: @@ -789,11 +789,11 @@ def _test_contact_position(shape: RigidObject, sensor: ContactSensor, mode: Cont # check shape of the contact_pos_w tensor (wp.to_torch expands vec3f -> float32 trailing dim) num_bodies = sensor.num_bodies - contact_pos_torch = wp.to_torch(sensor._data.contact_pos_w) + contact_pos_torch = sensor._data.contact_pos_w.torch assert contact_pos_torch.shape == (sensor.num_instances // num_bodies, num_bodies, 1, 3) # check contact positions if mode == ContactTestMode.IN_CONTACT: - pos_w_torch = wp.to_torch(sensor._data.pos_w) + pos_w_torch = sensor._data.pos_w.torch contact_position = pos_w_torch + torch.tensor([[0.0, 0.0, -shape.cfg.spawn.radius]], device=pos_w_torch.device) assert torch.all( torch.abs(torch.linalg.norm(contact_pos_torch - contact_position.unsqueeze(1), ord=2, dim=-1)) < 1e-2 @@ -828,10 +828,10 @@ def _check_prim_contact_state_times( in_air = True if expected_contact_time > 0.0: in_contact = True - measured_contact_time = wp.to_torch(sensor.data.current_contact_time) - measured_air_time = wp.to_torch(sensor.data.current_air_time) - measured_last_contact_time = wp.to_torch(sensor.data.last_contact_time) - measured_last_air_time = wp.to_torch(sensor.data.last_air_time) + measured_contact_time = sensor.data.current_contact_time.torch + measured_air_time = sensor.data.current_air_time.torch + measured_last_contact_time = sensor.data.last_contact_time.torch + measured_last_air_time = sensor.data.last_air_time.torch # check current contact state assert pytest.approx(measured_contact_time.item(), 0.01) == expected_contact_time assert pytest.approx(measured_air_time.item(), 0.01) == expected_air_time @@ -839,8 +839,8 @@ def _check_prim_contact_state_times( assert pytest.approx(measured_last_contact_time.item(), 0.01) == expected_last_contact_time assert pytest.approx(measured_last_air_time.item(), 0.01) == expected_last_air_time # check current contact mode - assert wp.to_torch(sensor.compute_first_contact(dt=dt)).item() == in_contact - assert wp.to_torch(sensor.compute_first_air(dt=dt)).item() == in_air + assert sensor.compute_first_contact(dt=dt).torch.item() == in_contact + assert sensor.compute_first_air(dt=dt).torch.item() == in_air def _perform_sim_step(sim, scene, sim_dt): diff --git a/source/isaaclab_physx/test/sensors/test_frame_transformer.py b/source/isaaclab_physx/test/sensors/test_frame_transformer.py index 96c9eabec2c2..631e9ba118dd 100644 --- a/source/isaaclab_physx/test/sensors/test_frame_transformer.py +++ b/source/isaaclab_physx/test/sensors/test_frame_transformer.py @@ -17,7 +17,6 @@ import pytest import scipy.spatial.transform as tf import torch -import warp as wp import isaaclab.sim as sim_utils import isaaclab.utils.math as math_utils @@ -146,7 +145,7 @@ def test_frame_transformer_feet_wrt_base(sim): feet_indices = [feet_indices[i] for i in reordering_indices] # default joint targets - default_actions = wp.to_torch(scene.articulations["robot"].data.default_joint_pos).clone() + default_actions = scene.articulations["robot"].data.default_joint_pos.torch.clone() # Define simulation stepping sim_dt = sim.get_physics_dt() # Simulate physics @@ -154,10 +153,10 @@ def test_frame_transformer_feet_wrt_base(sim): # # reset if count % 25 == 0: # reset root state - root_state = wp.to_torch(scene.articulations["robot"].data.default_root_state).clone() + root_state = scene.articulations["robot"].data.default_root_state.torch.clone() root_state[:, :3] += scene.env_origins - joint_pos = scene.articulations["robot"].data.default_joint_pos - joint_vel = scene.articulations["robot"].data.default_joint_vel + joint_pos = scene.articulations["robot"].data.default_joint_pos.torch + joint_vel = scene.articulations["robot"].data.default_joint_vel.torch # -- set root state # -- robot scene.articulations["robot"].write_root_pose_to_sim(root_state[:, :7]) @@ -178,14 +177,14 @@ def test_frame_transformer_feet_wrt_base(sim): # check absolute frame transforms in world frame # -- ground-truth - root_pose_w = wp.to_torch(scene.articulations["robot"].data.root_pose_w) - feet_pos_w_gt = wp.to_torch(scene.articulations["robot"].data.body_pos_w)[:, feet_indices] - feet_quat_w_gt = wp.to_torch(scene.articulations["robot"].data.body_quat_w)[:, feet_indices] + root_pose_w = scene.articulations["robot"].data.root_pose_w.torch + feet_pos_w_gt = scene.articulations["robot"].data.body_pos_w.torch[:, feet_indices] + feet_quat_w_gt = scene.articulations["robot"].data.body_quat_w.torch[:, feet_indices] # -- frame transformer - source_pos_w_tf = wp.to_torch(scene.sensors["frame_transformer"].data.source_pos_w) - source_quat_w_tf = wp.to_torch(scene.sensors["frame_transformer"].data.source_quat_w) - feet_pos_w_tf = wp.to_torch(scene.sensors["frame_transformer"].data.target_pos_w) - feet_quat_w_tf = wp.to_torch(scene.sensors["frame_transformer"].data.target_quat_w) + source_pos_w_tf = scene.sensors["frame_transformer"].data.source_pos_w.torch + source_quat_w_tf = scene.sensors["frame_transformer"].data.source_quat_w.torch + feet_pos_w_tf = scene.sensors["frame_transformer"].data.target_pos_w.torch + feet_quat_w_tf = scene.sensors["frame_transformer"].data.target_quat_w.torch # check if they are same torch.testing.assert_close(root_pose_w[:, :3], source_pos_w_tf) @@ -194,8 +193,8 @@ def test_frame_transformer_feet_wrt_base(sim): torch.testing.assert_close(feet_quat_w_gt, feet_quat_w_tf) # check if relative transforms are same - feet_pos_source_tf = wp.to_torch(scene.sensors["frame_transformer"].data.target_pos_source) - feet_quat_source_tf = wp.to_torch(scene.sensors["frame_transformer"].data.target_quat_source) + feet_pos_source_tf = scene.sensors["frame_transformer"].data.target_pos_source.torch + feet_quat_source_tf = scene.sensors["frame_transformer"].data.target_quat_source.torch for index in range(len(feet_indices)): # ground-truth foot_pos_b, foot_quat_b = math_utils.subtract_frame_transforms( @@ -244,7 +243,7 @@ def test_frame_transformer_feet_wrt_thigh(sim): assert scene.sensors["frame_transformer"].data.target_frame_names == user_feet_names # default joint targets - default_actions = wp.to_torch(scene.articulations["robot"].data.default_joint_pos).clone() + default_actions = scene.articulations["robot"].data.default_joint_pos.torch.clone() # Define simulation stepping sim_dt = sim.get_physics_dt() # Simulate physics @@ -252,10 +251,10 @@ def test_frame_transformer_feet_wrt_thigh(sim): # # reset if count % 25 == 0: # reset root state - root_state = wp.to_torch(scene.articulations["robot"].data.default_root_state).clone() + root_state = scene.articulations["robot"].data.default_root_state.torch.clone() root_state[:, :3] += scene.env_origins - joint_pos = scene.articulations["robot"].data.default_joint_pos - joint_vel = scene.articulations["robot"].data.default_joint_vel + joint_pos = scene.articulations["robot"].data.default_joint_pos.torch + joint_vel = scene.articulations["robot"].data.default_joint_vel.torch # -- set root state # -- robot scene.articulations["robot"].write_root_pose_to_sim(root_state[:, :7]) @@ -276,14 +275,14 @@ def test_frame_transformer_feet_wrt_thigh(sim): # check absolute frame transforms in world frame # -- ground-truth - source_pose_w_gt = wp.to_torch(scene.articulations["robot"].data.body_state_w)[:, source_frame_index, :7] - feet_pos_w_gt = wp.to_torch(scene.articulations["robot"].data.body_pos_w)[:, feet_indices] - feet_quat_w_gt = wp.to_torch(scene.articulations["robot"].data.body_quat_w)[:, feet_indices] + source_pose_w_gt = scene.articulations["robot"].data.body_state_w.torch[:, source_frame_index, :7] + feet_pos_w_gt = scene.articulations["robot"].data.body_pos_w.torch[:, feet_indices] + feet_quat_w_gt = scene.articulations["robot"].data.body_quat_w.torch[:, feet_indices] # -- frame transformer - source_pos_w_tf = wp.to_torch(scene.sensors["frame_transformer"].data.source_pos_w) - source_quat_w_tf = wp.to_torch(scene.sensors["frame_transformer"].data.source_quat_w) - feet_pos_w_tf = wp.to_torch(scene.sensors["frame_transformer"].data.target_pos_w) - feet_quat_w_tf = wp.to_torch(scene.sensors["frame_transformer"].data.target_quat_w) + source_pos_w_tf = scene.sensors["frame_transformer"].data.source_pos_w.torch + source_quat_w_tf = scene.sensors["frame_transformer"].data.source_quat_w.torch + feet_pos_w_tf = scene.sensors["frame_transformer"].data.target_pos_w.torch + feet_quat_w_tf = scene.sensors["frame_transformer"].data.target_quat_w.torch # check if they are same torch.testing.assert_close(source_pose_w_gt[:, :3], source_pos_w_tf) torch.testing.assert_close(source_pose_w_gt[:, 3:], source_quat_w_tf) @@ -291,8 +290,8 @@ def test_frame_transformer_feet_wrt_thigh(sim): torch.testing.assert_close(feet_quat_w_gt, feet_quat_w_tf) # check if relative transforms are same - feet_pos_source_tf = wp.to_torch(scene.sensors["frame_transformer"].data.target_pos_source) - feet_quat_source_tf = wp.to_torch(scene.sensors["frame_transformer"].data.target_quat_source) + feet_pos_source_tf = scene.sensors["frame_transformer"].data.target_pos_source.torch + feet_quat_source_tf = scene.sensors["frame_transformer"].data.target_quat_source.torch for index in range(len(feet_indices)): # ground-truth foot_pos_b, foot_quat_b = math_utils.subtract_frame_transforms( @@ -322,7 +321,7 @@ def test_frame_transformer_robot_body_to_external_cube(sim): sim.reset() # default joint targets - default_actions = wp.to_torch(scene.articulations["robot"].data.default_joint_pos).clone() + default_actions = scene.articulations["robot"].data.default_joint_pos.torch.clone() # Define simulation stepping sim_dt = sim.get_physics_dt() # Simulate physics @@ -330,10 +329,10 @@ def test_frame_transformer_robot_body_to_external_cube(sim): # # reset if count % 25 == 0: # reset root state - root_state = wp.to_torch(scene.articulations["robot"].data.default_root_state).clone() + root_state = scene.articulations["robot"].data.default_root_state.torch.clone() root_state[:, :3] += scene.env_origins - joint_pos = scene.articulations["robot"].data.default_joint_pos - joint_vel = scene.articulations["robot"].data.default_joint_vel + joint_pos = scene.articulations["robot"].data.default_joint_pos.torch + joint_vel = scene.articulations["robot"].data.default_joint_vel.torch # -- set root state # -- robot scene.articulations["robot"].write_root_pose_to_sim(root_state[:, :7]) @@ -354,14 +353,14 @@ def test_frame_transformer_robot_body_to_external_cube(sim): # check absolute frame transforms in world frame # -- ground-truth - root_pose_w = wp.to_torch(scene.articulations["robot"].data.root_pose_w) - cube_pos_w_gt = wp.to_torch(scene.rigid_objects["cube"].data.root_pos_w) - cube_quat_w_gt = wp.to_torch(scene.rigid_objects["cube"].data.root_quat_w) + root_pose_w = scene.articulations["robot"].data.root_pose_w.torch + cube_pos_w_gt = scene.rigid_objects["cube"].data.root_pos_w.torch + cube_quat_w_gt = scene.rigid_objects["cube"].data.root_quat_w.torch # -- frame transformer - source_pos_w_tf = wp.to_torch(scene.sensors["frame_transformer"].data.source_pos_w) - source_quat_w_tf = wp.to_torch(scene.sensors["frame_transformer"].data.source_quat_w) - cube_pos_w_tf = wp.to_torch(scene.sensors["frame_transformer"].data.target_pos_w).squeeze() - cube_quat_w_tf = wp.to_torch(scene.sensors["frame_transformer"].data.target_quat_w).squeeze() + source_pos_w_tf = scene.sensors["frame_transformer"].data.source_pos_w.torch + source_quat_w_tf = scene.sensors["frame_transformer"].data.source_quat_w.torch + cube_pos_w_tf = scene.sensors["frame_transformer"].data.target_pos_w.torch.squeeze() + cube_quat_w_tf = scene.sensors["frame_transformer"].data.target_quat_w.torch.squeeze() # check if they are same torch.testing.assert_close(root_pose_w[:, :3], source_pos_w_tf) @@ -370,8 +369,8 @@ def test_frame_transformer_robot_body_to_external_cube(sim): torch.testing.assert_close(cube_quat_w_gt, cube_quat_w_tf) # check if relative transforms are same - cube_pos_source_tf = wp.to_torch(scene.sensors["frame_transformer"].data.target_pos_source) - cube_quat_source_tf = wp.to_torch(scene.sensors["frame_transformer"].data.target_quat_source) + cube_pos_source_tf = scene.sensors["frame_transformer"].data.target_pos_source.torch + cube_quat_source_tf = scene.sensors["frame_transformer"].data.target_quat_source.torch # ground-truth cube_pos_b, cube_quat_b = math_utils.subtract_frame_transforms( root_pose_w[:, :3], root_pose_w[:, 3:], cube_pos_w_tf, cube_quat_w_tf @@ -426,7 +425,7 @@ def test_frame_transformer_offset_frames(sim): # # reset if count % 25 == 0: # reset root state - root_state = wp.to_torch(scene["cube"].data.default_root_state).clone() + root_state = scene["cube"].data.default_root_state.torch.clone() root_state[:, :3] += scene.env_origins # -- set root state # -- cube @@ -444,13 +443,13 @@ def test_frame_transformer_offset_frames(sim): # check absolute frame transforms in world frame # -- ground-truth - cube_pos_w_gt = wp.to_torch(scene["cube"].data.root_pos_w) - cube_quat_w_gt = wp.to_torch(scene["cube"].data.root_quat_w) + cube_pos_w_gt = scene["cube"].data.root_pos_w.torch + cube_quat_w_gt = scene["cube"].data.root_quat_w.torch # -- frame transformer - source_pos_w_tf = wp.to_torch(scene.sensors["frame_transformer"].data.source_pos_w) - source_quat_w_tf = wp.to_torch(scene.sensors["frame_transformer"].data.source_quat_w) - target_pos_w_tf = wp.to_torch(scene.sensors["frame_transformer"].data.target_pos_w).squeeze() - target_quat_w_tf = wp.to_torch(scene.sensors["frame_transformer"].data.target_quat_w).squeeze() + source_pos_w_tf = scene.sensors["frame_transformer"].data.source_pos_w.torch + source_quat_w_tf = scene.sensors["frame_transformer"].data.source_quat_w.torch + target_pos_w_tf = scene.sensors["frame_transformer"].data.target_pos_w.torch.squeeze() + target_quat_w_tf = scene.sensors["frame_transformer"].data.target_quat_w.torch.squeeze() target_frame_names = scene.sensors["frame_transformer"].data.target_frame_names cube_center_idx = target_frame_names.index("CUBE_CENTER") @@ -506,7 +505,7 @@ def test_frame_transformer_all_bodies(sim): reordering_indices = [target_frame_names.index(name) for name in articulation_body_names] # default joint targets - default_actions = wp.to_torch(scene.articulations["robot"].data.default_joint_pos).clone() + default_actions = scene.articulations["robot"].data.default_joint_pos.torch.clone() # Define simulation stepping sim_dt = sim.get_physics_dt() # Simulate physics @@ -514,10 +513,10 @@ def test_frame_transformer_all_bodies(sim): # # reset if count % 25 == 0: # reset root state - root_state = wp.to_torch(scene.articulations["robot"].data.default_root_state).clone() + root_state = scene.articulations["robot"].data.default_root_state.torch.clone() root_state[:, :3] += scene.env_origins - joint_pos = scene.articulations["robot"].data.default_joint_pos - joint_vel = scene.articulations["robot"].data.default_joint_vel + joint_pos = scene.articulations["robot"].data.default_joint_pos.torch + joint_vel = scene.articulations["robot"].data.default_joint_vel.torch # -- set root state # -- robot scene.articulations["robot"].write_root_pose_to_sim(root_state[:, :7]) @@ -538,15 +537,15 @@ def test_frame_transformer_all_bodies(sim): # check absolute frame transforms in world frame # -- ground-truth - root_pose_w = wp.to_torch(scene.articulations["robot"].data.root_pose_w) - bodies_pos_w_gt = wp.to_torch(scene.articulations["robot"].data.body_pos_w) - bodies_quat_w_gt = wp.to_torch(scene.articulations["robot"].data.body_quat_w) + root_pose_w = scene.articulations["robot"].data.root_pose_w.torch + bodies_pos_w_gt = scene.articulations["robot"].data.body_pos_w.torch + bodies_quat_w_gt = scene.articulations["robot"].data.body_quat_w.torch # -- frame transformer - source_pos_w_tf = wp.to_torch(scene.sensors["frame_transformer"].data.source_pos_w) - source_quat_w_tf = wp.to_torch(scene.sensors["frame_transformer"].data.source_quat_w) - bodies_pos_w_tf = wp.to_torch(scene.sensors["frame_transformer"].data.target_pos_w) - bodies_quat_w_tf = wp.to_torch(scene.sensors["frame_transformer"].data.target_quat_w) + source_pos_w_tf = scene.sensors["frame_transformer"].data.source_pos_w.torch + source_quat_w_tf = scene.sensors["frame_transformer"].data.source_quat_w.torch + bodies_pos_w_tf = scene.sensors["frame_transformer"].data.target_pos_w.torch + bodies_quat_w_tf = scene.sensors["frame_transformer"].data.target_quat_w.torch # check if they are same torch.testing.assert_close(root_pose_w[:, :3], source_pos_w_tf) @@ -554,8 +553,8 @@ def test_frame_transformer_all_bodies(sim): torch.testing.assert_close(bodies_pos_w_gt, bodies_pos_w_tf[:, reordering_indices]) torch.testing.assert_close(bodies_quat_w_gt, bodies_quat_w_tf[:, reordering_indices]) - bodies_pos_source_tf = wp.to_torch(scene.sensors["frame_transformer"].data.target_pos_source) - bodies_quat_source_tf = wp.to_torch(scene.sensors["frame_transformer"].data.target_quat_source) + bodies_pos_source_tf = scene.sensors["frame_transformer"].data.target_pos_source.torch + bodies_quat_source_tf = scene.sensors["frame_transformer"].data.target_quat_source.torch # Go through each body and check if relative transforms are same for index in range(len(articulation_body_names)): @@ -709,22 +708,22 @@ class MultiRobotSceneCfg(InteractiveSceneCfg): # Reset periodically if count % 10 == 0: # Reset robot - root_state = wp.to_torch(scene.articulations["robot"].data.default_root_state).clone() + root_state = scene.articulations["robot"].data.default_root_state.torch.clone() root_state[:, :3] += scene.env_origins scene.articulations["robot"].write_root_pose_to_sim(root_state[:, :7]) scene.articulations["robot"].write_root_velocity_to_sim(root_state[:, 7:]) scene.articulations["robot"].write_joint_state_to_sim( - scene.articulations["robot"].data.default_joint_pos, - scene.articulations["robot"].data.default_joint_vel, + scene.articulations["robot"].data.default_joint_pos.torch, + scene.articulations["robot"].data.default_joint_vel.torch, ) # Reset robot_1 - root_state_1 = wp.to_torch(scene.articulations["robot_1"].data.default_root_state).clone() + root_state_1 = scene.articulations["robot_1"].data.default_root_state.torch.clone() root_state_1[:, :3] += scene.env_origins scene.articulations["robot_1"].write_root_pose_to_sim(root_state_1[:, :7]) scene.articulations["robot_1"].write_root_velocity_to_sim(root_state_1[:, 7:]) scene.articulations["robot_1"].write_joint_state_to_sim( - scene.articulations["robot_1"].data.default_joint_pos, - scene.articulations["robot_1"].data.default_joint_vel, + scene.articulations["robot_1"].data.default_joint_pos.torch, + scene.articulations["robot_1"].data.default_joint_vel.torch, ) scene.reset() @@ -737,21 +736,21 @@ class MultiRobotSceneCfg(InteractiveSceneCfg): # Get frame transformer data frame_transformer_data = scene.sensors["frame_transformer"].data - source_pos_w = wp.to_torch(frame_transformer_data.source_pos_w) - source_quat_w = wp.to_torch(frame_transformer_data.source_quat_w) - target_pos_w = wp.to_torch(frame_transformer_data.target_pos_w) + source_pos_w = frame_transformer_data.source_pos_w.torch + source_quat_w = frame_transformer_data.source_quat_w.torch + target_pos_w = frame_transformer_data.target_pos_w.torch # Get ground truth positions and orientations (after scene.update() so they're current) - robot_lf_pos_w = wp.to_torch(scene.articulations["robot"].data.body_pos_w)[:, robot_lf_shank_body_idx] - robot_1_lf_pos_w = wp.to_torch(scene.articulations["robot_1"].data.body_pos_w)[:, robot_1_lf_shank_body_idx] - robot_rf_pos_w = wp.to_torch(scene.articulations["robot"].data.body_pos_w)[:, robot_rf_shank_body_idx] - robot_1_rf_pos_w = wp.to_torch(scene.articulations["robot_1"].data.body_pos_w)[:, robot_1_rf_shank_body_idx] + robot_lf_pos_w = scene.articulations["robot"].data.body_pos_w.torch[:, robot_lf_shank_body_idx] + robot_1_lf_pos_w = scene.articulations["robot_1"].data.body_pos_w.torch[:, robot_1_lf_shank_body_idx] + robot_rf_pos_w = scene.articulations["robot"].data.body_pos_w.torch[:, robot_rf_shank_body_idx] + robot_1_rf_pos_w = scene.articulations["robot_1"].data.body_pos_w.torch[:, robot_1_rf_shank_body_idx] # Get expected source frame positions and orientations (after scene.update() so they're current) - expected_source_base_pos_w = wp.to_torch(scene.articulations[expected_source_robot].data.body_pos_w)[ + expected_source_base_pos_w = scene.articulations[expected_source_robot].data.body_pos_w.torch[ :, expected_source_base_body_idx ] - expected_source_base_quat_w = wp.to_torch(scene.articulations[expected_source_robot].data.body_quat_w)[ + expected_source_base_quat_w = scene.articulations[expected_source_robot].data.body_quat_w.torch[ :, expected_source_base_body_idx ] diff --git a/source/isaaclab_physx/test/sensors/test_imu.py b/source/isaaclab_physx/test/sensors/test_imu.py index 7490b37eaf3f..641c72a8d4ba 100644 --- a/source/isaaclab_physx/test/sensors/test_imu.py +++ b/source/isaaclab_physx/test/sensors/test_imu.py @@ -247,22 +247,22 @@ def test_constant_velocity(setup_sim): if idx > 1: # check the imu accelerations torch.testing.assert_close( - wp.to_torch(scene.sensors["imu_ball"].data.lin_acc_b), + scene.sensors["imu_ball"].data.lin_acc_b.torch, prev_lin_acc_ball, rtol=1e-3, atol=1e-3, ) torch.testing.assert_close( - wp.to_torch(scene.sensors["imu_cube"].data.lin_acc_b), + scene.sensors["imu_cube"].data.lin_acc_b.torch, prev_lin_acc_cube, rtol=1e-3, atol=1e-3, ) # update previous values - prev_lin_acc_ball = wp.to_torch(scene.sensors["imu_ball"].data.lin_acc_b).clone() - prev_lin_acc_cube = wp.to_torch(scene.sensors["imu_cube"].data.lin_acc_b).clone() + prev_lin_acc_ball = scene.sensors["imu_ball"].data.lin_acc_b.torch.clone() + prev_lin_acc_cube = scene.sensors["imu_cube"].data.lin_acc_b.torch.clone() @pytest.mark.isaacsim_ci @@ -290,9 +290,9 @@ def test_constant_acceleration(setup_sim): # check the imu linear acceleration data (includes gravity since IMU always measures it) torch.testing.assert_close( - wp.to_torch(scene.sensors["imu_ball"].data.lin_acc_b), + scene.sensors["imu_ball"].data.lin_acc_b.torch, math_utils.quat_apply_inverse( - wp.to_torch(scene.rigid_objects["balls"].data.root_quat_w), + scene.rigid_objects["balls"].data.root_quat_w.torch, torch.tensor([[0.1, 0.0, 0.0]], dtype=torch.float32, device=scene.device).repeat(scene.num_envs, 1) / sim.get_physics_dt() + torch.tensor([[0.0, 0.0, 9.81]], dtype=torch.float32, device=scene.device).repeat(scene.num_envs, 1), @@ -303,8 +303,8 @@ def test_constant_acceleration(setup_sim): # check the angular velocity torch.testing.assert_close( - wp.to_torch(scene.sensors["imu_ball"].data.ang_vel_b), - wp.to_torch(scene.rigid_objects["balls"].data.root_ang_vel_b), + scene.sensors["imu_ball"].data.ang_vel_b.torch, + scene.rigid_objects["balls"].data.root_ang_vel_b.torch, rtol=1e-4, atol=1e-4, ) @@ -333,16 +333,16 @@ def test_single_dof_pendulum(setup_sim): # check the angular velocities of the imus between offset and direct definition torch.testing.assert_close( - wp.to_torch(base_data.ang_vel_b), - wp.to_torch(imu_data.ang_vel_b), + base_data.ang_vel_b.torch, + imu_data.ang_vel_b.torch, rtol=1e-4, atol=1e-4, ) # check the linear acceleration of the imus between offset and direct definition torch.testing.assert_close( - wp.to_torch(base_data.lin_acc_b), - wp.to_torch(imu_data.lin_acc_b), + base_data.lin_acc_b.torch, + imu_data.lin_acc_b.torch, rtol=1e-1, atol=1e-1, ) @@ -385,16 +385,16 @@ def test_indirect_attachment(setup_sim): # check the angular velocities of the imus between offset and direct definition torch.testing.assert_close( - wp.to_torch(base_data.ang_vel_b), - wp.to_torch(imu_data.ang_vel_b), + base_data.ang_vel_b.torch, + imu_data.ang_vel_b.torch, rtol=1e-4, atol=1e-4, ) # check the linear acceleration of the imus between offset and direct definition torch.testing.assert_close( - wp.to_torch(base_data.lin_acc_b), - wp.to_torch(imu_data.lin_acc_b), + base_data.lin_acc_b.torch, + imu_data.lin_acc_b.torch, rtol=1e-1, atol=1e-1, ) @@ -427,16 +427,16 @@ def test_offset_calculation(setup_sim): # check the linear accelerations torch.testing.assert_close( - wp.to_torch(scene.sensors["imu_robot_base"].data.lin_acc_b), - wp.to_torch(scene.sensors["imu_robot_imu_link"].data.lin_acc_b), + scene.sensors["imu_robot_base"].data.lin_acc_b.torch, + scene.sensors["imu_robot_imu_link"].data.lin_acc_b.torch, rtol=1e-4, atol=1e-4, ) # check the angular velocities torch.testing.assert_close( - wp.to_torch(scene.sensors["imu_robot_base"].data.ang_vel_b), - wp.to_torch(scene.sensors["imu_robot_imu_link"].data.ang_vel_b), + scene.sensors["imu_robot_base"].data.ang_vel_b.torch, + scene.sensors["imu_robot_imu_link"].data.ang_vel_b.torch, rtol=1e-4, atol=1e-4, ) diff --git a/source/isaaclab_physx/test/sensors/test_pva.py b/source/isaaclab_physx/test/sensors/test_pva.py index 624c1b279530..feb527f88c62 100644 --- a/source/isaaclab_physx/test/sensors/test_pva.py +++ b/source/isaaclab_physx/test/sensors/test_pva.py @@ -258,26 +258,26 @@ def test_constant_velocity(setup_sim): if idx > 1: # check the pva accelerations torch.testing.assert_close( - wp.to_torch(scene.sensors["pva_ball"].data.lin_acc_b), + scene.sensors["pva_ball"].data.lin_acc_b.torch, prev_lin_acc_ball, rtol=1e-3, atol=1e-3, ) torch.testing.assert_close( - wp.to_torch(scene.sensors["pva_ball"].data.ang_acc_b), + scene.sensors["pva_ball"].data.ang_acc_b.torch, prev_ang_acc_ball, rtol=1e-3, atol=1e-3, ) torch.testing.assert_close( - wp.to_torch(scene.sensors["pva_cube"].data.lin_acc_b), + scene.sensors["pva_cube"].data.lin_acc_b.torch, prev_lin_acc_cube, rtol=1e-3, atol=1e-3, ) torch.testing.assert_close( - wp.to_torch(scene.sensors["pva_cube"].data.ang_acc_b), + scene.sensors["pva_cube"].data.ang_acc_b.torch, prev_ang_acc_cube, rtol=1e-3, atol=1e-3, @@ -288,7 +288,7 @@ def test_constant_velocity(setup_sim): # setting v_0 (initial velocity) and then a calculation step of v_i = v_0 + a*dt. Consequently, # the data.lin_vel_b is returning approx. v_i. torch.testing.assert_close( - wp.to_torch(scene.sensors["pva_ball"].data.lin_vel_b), + scene.sensors["pva_ball"].data.lin_vel_b.torch, torch.tensor([[1.0, 0.0, -scene.physics_dt * 9.81]], dtype=torch.float32, device=scene.device).repeat( scene.num_envs, 1 ), @@ -296,7 +296,7 @@ def test_constant_velocity(setup_sim): atol=1e-4, ) torch.testing.assert_close( - wp.to_torch(scene.sensors["pva_cube"].data.lin_vel_b), + scene.sensors["pva_cube"].data.lin_vel_b.torch, torch.tensor([[1.0, 0.0, -scene.physics_dt * 9.81]], dtype=torch.float32, device=scene.device).repeat( scene.num_envs, 1 ), @@ -305,10 +305,10 @@ def test_constant_velocity(setup_sim): ) # update previous values - prev_lin_acc_ball = wp.to_torch(scene.sensors["pva_ball"].data.lin_acc_b).clone() - prev_ang_acc_ball = wp.to_torch(scene.sensors["pva_ball"].data.ang_acc_b).clone() - prev_lin_acc_cube = wp.to_torch(scene.sensors["pva_cube"].data.lin_acc_b).clone() - prev_ang_acc_cube = wp.to_torch(scene.sensors["pva_cube"].data.ang_acc_b).clone() + prev_lin_acc_ball = scene.sensors["pva_ball"].data.lin_acc_b.torch.clone() + prev_ang_acc_ball = scene.sensors["pva_ball"].data.ang_acc_b.torch.clone() + prev_lin_acc_cube = scene.sensors["pva_cube"].data.lin_acc_b.torch.clone() + prev_ang_acc_cube = scene.sensors["pva_cube"].data.ang_acc_b.torch.clone() @pytest.mark.isaacsim_ci @@ -336,9 +336,9 @@ def test_constant_acceleration(setup_sim): # check the pva data torch.testing.assert_close( - wp.to_torch(scene.sensors["pva_ball"].data.lin_acc_b), + scene.sensors["pva_ball"].data.lin_acc_b.torch, math_utils.quat_apply_inverse( - wp.to_torch(scene.rigid_objects["balls"].data.root_quat_w), + scene.rigid_objects["balls"].data.root_quat_w.torch, torch.tensor([[0.1, 0.0, 0.0]], dtype=torch.float32, device=scene.device).repeat(scene.num_envs, 1) / sim.get_physics_dt(), ), @@ -348,8 +348,8 @@ def test_constant_acceleration(setup_sim): # check the angular velocity torch.testing.assert_close( - wp.to_torch(scene.sensors["pva_ball"].data.ang_vel_b), - wp.to_torch(scene.rigid_objects["balls"].data.root_ang_vel_b), + scene.sensors["pva_ball"].data.ang_vel_b.torch, + scene.rigid_objects["balls"].data.root_ang_vel_b.torch, rtol=1e-4, atol=1e-4, ) @@ -372,25 +372,21 @@ def test_single_dof_pendulum(setup_sim): scene.update(sim.get_physics_dt()) # get pendulum joint state - joint_pos = wp.to_torch(scene.articulations["pendulum"].data.joint_pos) - joint_vel = wp.to_torch(scene.articulations["pendulum"].data.joint_vel) - joint_acc = wp.to_torch(scene.articulations["pendulum"].data.joint_acc) + joint_pos = scene.articulations["pendulum"].data.joint_pos.torch + joint_vel = scene.articulations["pendulum"].data.joint_vel.torch + joint_acc = scene.articulations["pendulum"].data.joint_acc.torch # PVA and base data pva_data = scene.sensors["pva_pendulum_imu_link"].data base_data = scene.sensors["pva_pendulum_base"].data # extract imu_link pva_sensor dynamics - lin_vel_w_imu_link = math_utils.quat_apply(wp.to_torch(pva_data.quat_w), wp.to_torch(pva_data.lin_vel_b)) - lin_acc_w_imu_link = math_utils.quat_apply(wp.to_torch(pva_data.quat_w), wp.to_torch(pva_data.lin_acc_b)) + lin_vel_w_imu_link = math_utils.quat_apply(pva_data.quat_w.torch, pva_data.lin_vel_b.torch) + lin_acc_w_imu_link = math_utils.quat_apply(pva_data.quat_w.torch, pva_data.lin_acc_b.torch) # calculate the joint dynamics from the pva_sensor (y axis of imu_link is parallel to joint axis of pendulum) - joint_vel_pva = math_utils.quat_apply(wp.to_torch(pva_data.quat_w), wp.to_torch(pva_data.ang_vel_b))[ - ..., 1 - ].unsqueeze(-1) - joint_acc_pva = math_utils.quat_apply(wp.to_torch(pva_data.quat_w), wp.to_torch(pva_data.ang_acc_b))[ - ..., 1 - ].unsqueeze(-1) + joint_vel_pva = math_utils.quat_apply(pva_data.quat_w.torch, pva_data.ang_vel_b.torch)[..., 1].unsqueeze(-1) + joint_acc_pva = math_utils.quat_apply(pva_data.quat_w.torch, pva_data.ang_acc_b.torch)[..., 1].unsqueeze(-1) # calculate analytical solution vx = -joint_vel * pend_length * torch.sin(joint_pos) @@ -409,9 +405,9 @@ def test_single_dof_pendulum(setup_sim): # compare pva projected gravity gravity_dir_w = torch.tensor((0.0, 0.0, -1.0), device=scene.device).repeat(2, 1) - gravity_dir_b = math_utils.quat_apply_inverse(wp.to_torch(pva_data.quat_w), gravity_dir_w) + gravity_dir_b = math_utils.quat_apply_inverse(pva_data.quat_w.torch, gravity_dir_w) torch.testing.assert_close( - wp.to_torch(pva_data.projected_gravity_b), + pva_data.projected_gravity_b.torch, gravity_dir_b, ) @@ -446,47 +442,47 @@ def test_single_dof_pendulum(setup_sim): # check the position between offset and pva definition torch.testing.assert_close( - wp.to_torch(base_data.pos_w), - wp.to_torch(pva_data.pos_w), + base_data.pos_w.torch, + pva_data.pos_w.torch, rtol=1e-5, atol=1e-5, ) # check the orientation between offset and pva definition torch.testing.assert_close( - wp.to_torch(base_data.quat_w), - wp.to_torch(pva_data.quat_w), + base_data.quat_w.torch, + pva_data.quat_w.torch, rtol=1e-4, atol=1e-4, ) # check the angular velocities of the pvas between offset and pva definition torch.testing.assert_close( - wp.to_torch(base_data.ang_vel_b), - wp.to_torch(pva_data.ang_vel_b), + base_data.ang_vel_b.torch, + pva_data.ang_vel_b.torch, rtol=1e-4, atol=1e-4, ) # check the angular acceleration of the pvas between offset and pva definition torch.testing.assert_close( - wp.to_torch(base_data.ang_acc_b), - wp.to_torch(pva_data.ang_acc_b), + base_data.ang_acc_b.torch, + pva_data.ang_acc_b.torch, rtol=1e-4, atol=1e-4, ) # check the linear velocity of the pvas between offset and pva definition torch.testing.assert_close( - wp.to_torch(base_data.lin_vel_b), - wp.to_torch(pva_data.lin_vel_b), + base_data.lin_vel_b.torch, + pva_data.lin_vel_b.torch, rtol=1e-2, atol=5e-3, ) # check the linear acceleration of the pvas between offset and pva definition torch.testing.assert_close( - wp.to_torch(base_data.lin_acc_b), - wp.to_torch(pva_data.lin_acc_b), + base_data.lin_acc_b.torch, + pva_data.lin_acc_b.torch, rtol=1e-1, atol=1e-1, ) @@ -509,9 +505,9 @@ def test_indirect_attachment(setup_sim): scene.update(sim.get_physics_dt()) # get pendulum joint state - joint_pos = wp.to_torch(scene.articulations["pendulum2"].data.joint_pos) - joint_vel = wp.to_torch(scene.articulations["pendulum2"].data.joint_vel) - joint_acc = wp.to_torch(scene.articulations["pendulum2"].data.joint_acc) + joint_pos = scene.articulations["pendulum2"].data.joint_pos.torch + joint_vel = scene.articulations["pendulum2"].data.joint_vel.torch + joint_acc = scene.articulations["pendulum2"].data.joint_acc.torch pva = scene.sensors["pva_indirect_pendulum_link"] pva_base = scene.sensors["pva_indirect_pendulum_base"] @@ -531,16 +527,12 @@ def test_indirect_attachment(setup_sim): pva_data = scene.sensors["pva_indirect_pendulum_link"].data base_data = scene.sensors["pva_indirect_pendulum_base"].data # extract imu_link pva_sensor dynamics - lin_vel_w_imu_link = math_utils.quat_apply(wp.to_torch(pva_data.quat_w), wp.to_torch(pva_data.lin_vel_b)) - lin_acc_w_imu_link = math_utils.quat_apply(wp.to_torch(pva_data.quat_w), wp.to_torch(pva_data.lin_acc_b)) + lin_vel_w_imu_link = math_utils.quat_apply(pva_data.quat_w.torch, pva_data.lin_vel_b.torch) + lin_acc_w_imu_link = math_utils.quat_apply(pva_data.quat_w.torch, pva_data.lin_acc_b.torch) # calculate the joint dynamics from the pva_sensor (y axis of imu_link is parallel to joint axis of pendulum) - joint_vel_pva = math_utils.quat_apply(wp.to_torch(pva_data.quat_w), wp.to_torch(pva_data.ang_vel_b))[ - ..., 1 - ].unsqueeze(-1) - joint_acc_pva = math_utils.quat_apply(wp.to_torch(pva_data.quat_w), wp.to_torch(pva_data.ang_acc_b))[ - ..., 1 - ].unsqueeze(-1) + joint_vel_pva = math_utils.quat_apply(pva_data.quat_w.torch, pva_data.ang_vel_b.torch)[..., 1].unsqueeze(-1) + joint_acc_pva = math_utils.quat_apply(pva_data.quat_w.torch, pva_data.ang_acc_b.torch)[..., 1].unsqueeze(-1) # calculate analytical solution vx = -joint_vel * pend_length * torch.sin(joint_pos) @@ -559,9 +551,9 @@ def test_indirect_attachment(setup_sim): # compare pva projected gravity gravity_dir_w = torch.tensor((0.0, 0.0, -1.0), device=scene.device).repeat(2, 1) - gravity_dir_b = math_utils.quat_apply_inverse(wp.to_torch(pva_data.quat_w), gravity_dir_w) + gravity_dir_b = math_utils.quat_apply_inverse(pva_data.quat_w.torch, gravity_dir_w) torch.testing.assert_close( - wp.to_torch(pva_data.projected_gravity_b), + pva_data.projected_gravity_b.torch, gravity_dir_b, ) @@ -596,47 +588,47 @@ def test_indirect_attachment(setup_sim): # check the position between offset and pva definition torch.testing.assert_close( - wp.to_torch(base_data.pos_w), - wp.to_torch(pva_data.pos_w), + base_data.pos_w.torch, + pva_data.pos_w.torch, rtol=1e-5, atol=1e-5, ) # check the orientation between offset and pva definition torch.testing.assert_close( - wp.to_torch(base_data.quat_w), - wp.to_torch(pva_data.quat_w), + base_data.quat_w.torch, + pva_data.quat_w.torch, rtol=1e-4, atol=1e-4, ) # check the angular velocities of the pvas between offset and pva definition torch.testing.assert_close( - wp.to_torch(base_data.ang_vel_b), - wp.to_torch(pva_data.ang_vel_b), + base_data.ang_vel_b.torch, + pva_data.ang_vel_b.torch, rtol=1e-4, atol=1e-4, ) # check the angular acceleration of the pvas between offset and pva definition torch.testing.assert_close( - wp.to_torch(base_data.ang_acc_b), - wp.to_torch(pva_data.ang_acc_b), + base_data.ang_acc_b.torch, + pva_data.ang_acc_b.torch, rtol=1e-4, atol=1e-4, ) # check the linear velocity of the pvas between offset and pva definition torch.testing.assert_close( - wp.to_torch(base_data.lin_vel_b), - wp.to_torch(pva_data.lin_vel_b), + base_data.lin_vel_b.torch, + pva_data.lin_vel_b.torch, rtol=1e-2, atol=5e-3, ) # check the linear acceleration of the pvas between offset and pva definition torch.testing.assert_close( - wp.to_torch(base_data.lin_acc_b), - wp.to_torch(pva_data.lin_acc_b), + base_data.lin_acc_b.torch, + pva_data.lin_acc_b.torch, rtol=1e-1, atol=1e-1, ) @@ -669,50 +661,50 @@ def test_offset_calculation(setup_sim): # check the accelerations torch.testing.assert_close( - wp.to_torch(scene.sensors["pva_robot_base"].data.lin_acc_b), - wp.to_torch(scene.sensors["pva_robot_imu_link"].data.lin_acc_b), + scene.sensors["pva_robot_base"].data.lin_acc_b.torch, + scene.sensors["pva_robot_imu_link"].data.lin_acc_b.torch, rtol=1e-4, atol=1e-4, ) torch.testing.assert_close( - wp.to_torch(scene.sensors["pva_robot_base"].data.ang_acc_b), - wp.to_torch(scene.sensors["pva_robot_imu_link"].data.ang_acc_b), + scene.sensors["pva_robot_base"].data.ang_acc_b.torch, + scene.sensors["pva_robot_imu_link"].data.ang_acc_b.torch, rtol=1e-4, atol=1e-4, ) # check the velocities torch.testing.assert_close( - wp.to_torch(scene.sensors["pva_robot_base"].data.ang_vel_b), - wp.to_torch(scene.sensors["pva_robot_imu_link"].data.ang_vel_b), + scene.sensors["pva_robot_base"].data.ang_vel_b.torch, + scene.sensors["pva_robot_imu_link"].data.ang_vel_b.torch, rtol=1e-4, atol=1e-4, ) torch.testing.assert_close( - wp.to_torch(scene.sensors["pva_robot_base"].data.lin_vel_b), - wp.to_torch(scene.sensors["pva_robot_imu_link"].data.lin_vel_b), + scene.sensors["pva_robot_base"].data.lin_vel_b.torch, + scene.sensors["pva_robot_imu_link"].data.lin_vel_b.torch, rtol=1e-4, atol=1e-4, ) # check the orientation torch.testing.assert_close( - wp.to_torch(scene.sensors["pva_robot_base"].data.quat_w), - wp.to_torch(scene.sensors["pva_robot_imu_link"].data.quat_w), + scene.sensors["pva_robot_base"].data.quat_w.torch, + scene.sensors["pva_robot_imu_link"].data.quat_w.torch, rtol=1e-4, atol=1e-4, ) # check the position torch.testing.assert_close( - wp.to_torch(scene.sensors["pva_robot_base"].data.pos_w), - wp.to_torch(scene.sensors["pva_robot_imu_link"].data.pos_w), + scene.sensors["pva_robot_base"].data.pos_w.torch, + scene.sensors["pva_robot_imu_link"].data.pos_w.torch, rtol=1e-4, atol=1e-4, ) # check the projected gravity torch.testing.assert_close( - wp.to_torch(scene.sensors["pva_robot_base"].data.projected_gravity_b), - wp.to_torch(scene.sensors["pva_robot_imu_link"].data.projected_gravity_b), + scene.sensors["pva_robot_base"].data.projected_gravity_b.torch, + scene.sensors["pva_robot_imu_link"].data.projected_gravity_b.torch, rtol=1e-4, atol=1e-4, ) diff --git a/source/isaaclab_tasks/config/extension.toml b/source/isaaclab_tasks/config/extension.toml index 29bf11d4859a..158c271c5feb 100644 --- a/source/isaaclab_tasks/config/extension.toml +++ b/source/isaaclab_tasks/config/extension.toml @@ -1,7 +1,7 @@ [package] # Note: Semantic Versioning is used: https://semver.org/ -version = "1.5.28" +version = "1.5.29" # Description title = "Isaac Lab Environments" diff --git a/source/isaaclab_tasks/docs/CHANGELOG.rst b/source/isaaclab_tasks/docs/CHANGELOG.rst index 54502c8a8354..42a68f5ae361 100644 --- a/source/isaaclab_tasks/docs/CHANGELOG.rst +++ b/source/isaaclab_tasks/docs/CHANGELOG.rst @@ -1,7 +1,7 @@ Changelog --------- -1.5.28 (2026-04-24) +1.5.29 (2026-04-27) ~~~~~~~~~~~~~~~~~~~ Added @@ -31,7 +31,7 @@ Changed ``(-1.0, 3.0)`` additive overrides on A1/Go1/Go2. -1.5.27 (2026-04-25) +1.5.28 (2026-04-27) ~~~~~~~~~~~~~~~~~~~ Changed @@ -41,7 +41,7 @@ Changed ``isaacsim.core.experimental.utils.app.enable_extension`` (non-deprecated Isaac Sim path). -1.5.26 (2026-04-25) +1.5.27 (2026-04-27) ~~~~~~~~~~~~~~~~~~~ Changed @@ -56,7 +56,7 @@ Changed sequence instead of ``torch.tensor``, dropping the ``torch`` import from those modules. -1.5.25 (2026-04-24) +1.5.26 (2026-04-27) ~~~~~~~~~~~~~~~~~~~ Changed @@ -65,6 +65,17 @@ Changed * Migrated golden images ``source/isaaclab_tasks/test/golden_images/**/*.png`` to Git LFS. +1.5.25 (2026-04-23) +~~~~~~~~~~~~~~~~~~~ + +Changed +^^^^^^^ + +* Migrated all data property accesses from ``wp.to_torch(data.property)`` to + ``data.property.torch`` to match the new :class:`~isaaclab.utils.warp.ProxyArray` + return type introduced in ``isaaclab`` 4.6.13. + + 1.5.24 (2026-04-22) ~~~~~~~~~~~~~~~~~~~ diff --git a/source/isaaclab_tasks/isaaclab_tasks/direct/anymal_c/anymal_c_env.py b/source/isaaclab_tasks/isaaclab_tasks/direct/anymal_c/anymal_c_env.py index e58e377a0d63..e628d98e533a 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/direct/anymal_c/anymal_c_env.py +++ b/source/isaaclab_tasks/isaaclab_tasks/direct/anymal_c/anymal_c_env.py @@ -76,9 +76,7 @@ def _setup_scene(self): def _pre_physics_step(self, actions: torch.Tensor): self._actions = actions.clone() - self._processed_actions = self.cfg.action_scale * self._actions + wp.to_torch( - self._robot.data.default_joint_pos - ) + self._processed_actions = self.cfg.action_scale * self._actions + self._robot.data.default_joint_pos.torch def _apply_action(self): self._robot.set_joint_position_target_index(target=self._processed_actions) @@ -88,20 +86,20 @@ def _get_observations(self) -> dict: height_data = None if isinstance(self.cfg, AnymalCRoughEnvCfg): height_data = ( - wp.to_torch(self._height_scanner.data.pos_w)[:, 2].unsqueeze(1) - - wp.to_torch(self._height_scanner.data.ray_hits_w)[..., 2] + self._height_scanner.data.pos_w.torch[:, 2].unsqueeze(1) + - self._height_scanner.data.ray_hits_w.torch[..., 2] - 0.5 ).clip(-1.0, 1.0) obs = torch.cat( [ tensor for tensor in ( - wp.to_torch(self._robot.data.root_lin_vel_b), - wp.to_torch(self._robot.data.root_ang_vel_b), - wp.to_torch(self._robot.data.projected_gravity_b), + self._robot.data.root_lin_vel_b.torch, + self._robot.data.root_ang_vel_b.torch, + self._robot.data.projected_gravity_b.torch, self._commands, - wp.to_torch(self._robot.data.joint_pos) - wp.to_torch(self._robot.data.default_joint_pos), - wp.to_torch(self._robot.data.joint_vel), + self._robot.data.joint_pos.torch - self._robot.data.default_joint_pos.torch, + self._robot.data.joint_vel.torch, height_data, self._actions, ) @@ -115,37 +113,37 @@ def _get_observations(self) -> dict: def _get_rewards(self) -> torch.Tensor: # linear velocity tracking lin_vel_error = torch.sum( - torch.square(self._commands[:, :2] - wp.to_torch(self._robot.data.root_lin_vel_b)[:, :2]), dim=1 + torch.square(self._commands[:, :2] - self._robot.data.root_lin_vel_b.torch[:, :2]), dim=1 ) lin_vel_error_mapped = torch.exp(-lin_vel_error / 0.25) # yaw rate tracking - yaw_rate_error = torch.square(self._commands[:, 2] - wp.to_torch(self._robot.data.root_ang_vel_b)[:, 2]) + yaw_rate_error = torch.square(self._commands[:, 2] - self._robot.data.root_ang_vel_b.torch[:, 2]) yaw_rate_error_mapped = torch.exp(-yaw_rate_error / 0.25) # z velocity tracking - z_vel_error = torch.square(wp.to_torch(self._robot.data.root_lin_vel_b)[:, 2]) + z_vel_error = torch.square(self._robot.data.root_lin_vel_b.torch[:, 2]) # angular velocity x/y - ang_vel_error = torch.sum(torch.square(wp.to_torch(self._robot.data.root_ang_vel_b)[:, :2]), dim=1) + ang_vel_error = torch.sum(torch.square(self._robot.data.root_ang_vel_b.torch[:, :2]), dim=1) # joint torques - joint_torques = torch.sum(torch.square(wp.to_torch(self._robot.data.applied_torque)), dim=1) + joint_torques = torch.sum(torch.square(self._robot.data.applied_torque.torch), dim=1) # joint acceleration - joint_accel = torch.sum(torch.square(wp.to_torch(self._robot.data.joint_acc)), dim=1) + joint_accel = torch.sum(torch.square(self._robot.data.joint_acc.torch), dim=1) # action rate action_rate = torch.sum(torch.square(self._actions - self._previous_actions), dim=1) # feet air time - first_contact = wp.to_torch(self._contact_sensor.compute_first_contact(self.step_dt))[:, self._feet_ids] - last_air_time = wp.to_torch(self._contact_sensor.data.last_air_time)[:, self._feet_ids] + first_contact = self._contact_sensor.compute_first_contact(self.step_dt).torch[:, self._feet_ids] + last_air_time = self._contact_sensor.data.last_air_time.torch[:, self._feet_ids] air_time = torch.sum((last_air_time - 0.5) * first_contact, dim=1) * ( torch.linalg.norm(self._commands[:, :2], dim=1) > 0.1 ) # undesired contacts - net_contact_forces = wp.to_torch(self._contact_sensor.data.net_forces_w_history) + net_contact_forces = self._contact_sensor.data.net_forces_w_history.torch is_contact = ( torch.max(torch.linalg.norm(net_contact_forces[:, :, self._undesired_contact_body_ids], dim=-1), dim=1)[0] > 1.0 ) contacts = torch.sum(is_contact, dim=1) # flat orientation - flat_orientation = torch.sum(torch.square(wp.to_torch(self._robot.data.projected_gravity_b)[:, :2]), dim=1) + flat_orientation = torch.sum(torch.square(self._robot.data.projected_gravity_b.torch[:, :2]), dim=1) rewards = { "track_lin_vel_xy_exp": lin_vel_error_mapped * self.cfg.lin_vel_reward_scale * self.step_dt, @@ -167,7 +165,7 @@ def _get_rewards(self) -> torch.Tensor: def _get_dones(self) -> tuple[torch.Tensor, torch.Tensor]: time_out = self.episode_length_buf >= self.max_episode_length - 1 - net_contact_forces = wp.to_torch(self._contact_sensor.data.net_forces_w_history) + net_contact_forces = self._contact_sensor.data.net_forces_w_history.torch died = torch.any( torch.max(torch.linalg.norm(net_contact_forces[:, :, self._base_id], dim=-1), dim=1)[0] > 1.0, dim=1 ) @@ -186,10 +184,10 @@ def _reset_idx(self, env_ids: torch.Tensor | None): # Sample new commands self._commands[env_ids] = torch.zeros_like(self._commands[env_ids]).uniform_(-1.0, 1.0) # Reset robot state - joint_pos = wp.to_torch(self._robot.data.default_joint_pos)[env_ids] - joint_vel = wp.to_torch(self._robot.data.default_joint_vel)[env_ids] - default_root_pose = wp.to_torch(self._robot.data.default_root_pose)[env_ids] - default_root_vel = wp.to_torch(self._robot.data.default_root_vel)[env_ids] + joint_pos = self._robot.data.default_joint_pos.torch[env_ids] + joint_vel = self._robot.data.default_joint_vel.torch[env_ids] + default_root_pose = self._robot.data.default_root_pose.torch[env_ids] + default_root_vel = self._robot.data.default_root_vel.torch[env_ids] default_root_pose[:, :3] += self._terrain.env_origins[env_ids] self._robot.write_root_pose_to_sim_index(root_pose=default_root_pose, env_ids=env_ids) self._robot.write_root_velocity_to_sim_index(root_velocity=default_root_vel, env_ids=env_ids) diff --git a/source/isaaclab_tasks/isaaclab_tasks/direct/automate/assembly_env.py b/source/isaaclab_tasks/isaaclab_tasks/direct/automate/assembly_env.py index cdd39875486b..3fcbe19d25bf 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/direct/automate/assembly_env.py +++ b/source/isaaclab_tasks/isaaclab_tasks/direct/automate/assembly_env.py @@ -275,18 +275,18 @@ def _setup_scene(self): def _compute_intermediate_values(self, dt): """Get values computed from raw tensors. This includes adding noise.""" # TODO: A lot of these can probably only be set once? - self.fixed_pos = wp.to_torch(self._fixed_asset.data.root_pos_w) - self.scene.env_origins - self.fixed_quat = wp.to_torch(self._fixed_asset.data.root_quat_w) + self.fixed_pos = self._fixed_asset.data.root_pos_w.torch - self.scene.env_origins + self.fixed_quat = self._fixed_asset.data.root_quat_w.torch - self.held_pos = wp.to_torch(self._held_asset.data.root_pos_w) - self.scene.env_origins - self.held_quat = wp.to_torch(self._held_asset.data.root_quat_w) + self.held_pos = self._held_asset.data.root_pos_w.torch - self.scene.env_origins + self.held_quat = self._held_asset.data.root_quat_w.torch self.fingertip_midpoint_pos = ( - wp.to_torch(self._robot.data.body_pos_w)[:, self.fingertip_body_idx] - self.scene.env_origins + self._robot.data.body_pos_w.torch[:, self.fingertip_body_idx] - self.scene.env_origins ) - self.fingertip_midpoint_quat = wp.to_torch(self._robot.data.body_quat_w)[:, self.fingertip_body_idx] - self.fingertip_midpoint_linvel = wp.to_torch(self._robot.data.body_lin_vel_w)[:, self.fingertip_body_idx] - self.fingertip_midpoint_angvel = wp.to_torch(self._robot.data.body_ang_vel_w)[:, self.fingertip_body_idx] + self.fingertip_midpoint_quat = self._robot.data.body_quat_w.torch[:, self.fingertip_body_idx] + self.fingertip_midpoint_linvel = self._robot.data.body_lin_vel_w.torch[:, self.fingertip_body_idx] + self.fingertip_midpoint_angvel = self._robot.data.body_ang_vel_w.torch[:, self.fingertip_body_idx] jacobians = wp.to_torch(self._robot.root_view.get_jacobians()) @@ -294,8 +294,8 @@ def _compute_intermediate_values(self, dt): self.right_finger_jacobian = jacobians[:, self.right_finger_body_idx - 1, 0:6, 0:7] self.fingertip_midpoint_jacobian = (self.left_finger_jacobian + self.right_finger_jacobian) * 0.5 self.arm_mass_matrix = wp.to_torch(self._robot.root_view.get_generalized_mass_matrices())[:, 0:7, 0:7] - self.joint_pos = wp.to_torch(self._robot.data.joint_pos).clone() - self.joint_vel = wp.to_torch(self._robot.data.joint_vel).clone() + self.joint_pos = self._robot.data.joint_pos.torch.clone() + self.joint_vel = self._robot.data.joint_vel.torch.clone() # Compute pose of gripper goal and top of socket in socket frame self.gripper_goal_pos, self.gripper_goal_quat = combine_frame_transforms( @@ -649,7 +649,7 @@ def _reset_idx(self, env_ids): def _set_assets_to_default_pose(self, env_ids): """Move assets to default pose before randomization.""" held_state = torch.cat( - [wp.to_torch(self._held_asset.data.default_root_pose), wp.to_torch(self._held_asset.data.default_root_vel)], + [self._held_asset.data.default_root_pose.torch, self._held_asset.data.default_root_vel.torch], dim=-1, )[env_ids].clone() held_state[:, 0:3] += self.scene.env_origins[env_ids] @@ -660,8 +660,8 @@ def _set_assets_to_default_pose(self, env_ids): fixed_state = torch.cat( [ - wp.to_torch(self._fixed_asset.data.default_root_pose), - wp.to_torch(self._fixed_asset.data.default_root_vel), + self._fixed_asset.data.default_root_pose.torch, + self._fixed_asset.data.default_root_vel.torch, ], dim=-1, )[env_ids].clone() @@ -739,7 +739,7 @@ def set_pos_inverse_kinematics(self, env_ids): def _set_franka_to_default_pose(self, joints, env_ids): """Return Franka to its default joint position.""" gripper_width = self.gripper_open_width - joint_pos = wp.to_torch(self._robot.data.default_joint_pos)[env_ids] + joint_pos = self._robot.data.default_joint_pos.torch[env_ids] joint_pos[:, 7:] = gripper_width # MIMIC joint_pos[:, :7] = torch.tensor(joints, device=self.device)[None, :] joint_vel = torch.zeros_like(joint_pos) @@ -764,8 +764,8 @@ def randomize_fixed_initial_state(self, env_ids): # (1.) Randomize fixed asset pose. fixed_state = torch.cat( [ - wp.to_torch(self._fixed_asset.data.default_root_pose), - wp.to_torch(self._fixed_asset.data.default_root_vel), + self._fixed_asset.data.default_root_pose.torch, + self._fixed_asset.data.default_root_vel.torch, ], dim=-1, )[env_ids].clone() @@ -819,7 +819,7 @@ def randomize_held_initial_state(self, env_ids, pre_grasp): # Set plug pos to assembled state, but offset plug Z-coordinate by height of socket, # minus curriculum displacement held_state = torch.cat( - [wp.to_torch(self._held_asset.data.default_root_pose), wp.to_torch(self._held_asset.data.default_root_vel)], + [self._held_asset.data.default_root_pose.torch, self._held_asset.data.default_root_vel.torch], dim=-1, ).clone() held_state[env_ids, 0:3] = self.fixed_pos[env_ids].clone() + self.scene.env_origins[env_ids] diff --git a/source/isaaclab_tasks/isaaclab_tasks/direct/automate/disassembly_env.py b/source/isaaclab_tasks/isaaclab_tasks/direct/automate/disassembly_env.py index 73822e43e9b9..e64d7018ce01 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/direct/automate/disassembly_env.py +++ b/source/isaaclab_tasks/isaaclab_tasks/direct/automate/disassembly_env.py @@ -202,18 +202,18 @@ def _setup_scene(self): def _compute_intermediate_values(self, dt): """Get values computed from raw tensors. This includes adding noise.""" # TODO: A lot of these can probably only be set once? - self.fixed_pos = wp.to_torch(self._fixed_asset.data.root_pos_w) - self.scene.env_origins - self.fixed_quat = wp.to_torch(self._fixed_asset.data.root_quat_w) + self.fixed_pos = self._fixed_asset.data.root_pos_w.torch - self.scene.env_origins + self.fixed_quat = self._fixed_asset.data.root_quat_w.torch - self.held_pos = wp.to_torch(self._held_asset.data.root_pos_w) - self.scene.env_origins - self.held_quat = wp.to_torch(self._held_asset.data.root_quat_w) + self.held_pos = self._held_asset.data.root_pos_w.torch - self.scene.env_origins + self.held_quat = self._held_asset.data.root_quat_w.torch self.fingertip_midpoint_pos = ( - wp.to_torch(self._robot.data.body_pos_w)[:, self.fingertip_body_idx] - self.scene.env_origins + self._robot.data.body_pos_w.torch[:, self.fingertip_body_idx] - self.scene.env_origins ) - self.fingertip_midpoint_quat = wp.to_torch(self._robot.data.body_quat_w)[:, self.fingertip_body_idx] - self.fingertip_midpoint_linvel = wp.to_torch(self._robot.data.body_lin_vel_w)[:, self.fingertip_body_idx] - self.fingertip_midpoint_angvel = wp.to_torch(self._robot.data.body_ang_vel_w)[:, self.fingertip_body_idx] + self.fingertip_midpoint_quat = self._robot.data.body_quat_w.torch[:, self.fingertip_body_idx] + self.fingertip_midpoint_linvel = self._robot.data.body_lin_vel_w.torch[:, self.fingertip_body_idx] + self.fingertip_midpoint_angvel = self._robot.data.body_ang_vel_w.torch[:, self.fingertip_body_idx] jacobians = wp.to_torch(self._robot.root_view.get_jacobians()) @@ -221,8 +221,8 @@ def _compute_intermediate_values(self, dt): self.right_finger_jacobian = jacobians[:, self.right_finger_body_idx - 1, 0:6, 0:7] self.fingertip_midpoint_jacobian = (self.left_finger_jacobian + self.right_finger_jacobian) * 0.5 self.arm_mass_matrix = wp.to_torch(self._robot.root_view.get_generalized_mass_matrices())[:, 0:7, 0:7] - self.joint_pos = wp.to_torch(self._robot.data.joint_pos).clone() - self.joint_vel = wp.to_torch(self._robot.data.joint_vel).clone() + self.joint_pos = self._robot.data.joint_pos.torch.clone() + self.joint_vel = self._robot.data.joint_vel.torch.clone() # Compute pose of gripper goal and top of socket in socket frame self.gripper_goal_pos, self.gripper_goal_quat = combine_frame_transforms( @@ -480,7 +480,7 @@ def _reset_idx(self, env_ids): def _set_assets_to_default_pose(self, env_ids): """Move assets to default pose before randomization.""" held_state = torch.cat( - [wp.to_torch(self._held_asset.data.default_root_pose), wp.to_torch(self._held_asset.data.default_root_vel)], + [self._held_asset.data.default_root_pose.torch, self._held_asset.data.default_root_vel.torch], dim=-1, )[env_ids].clone() held_state[:, 0:3] += self.scene.env_origins[env_ids] @@ -491,8 +491,8 @@ def _set_assets_to_default_pose(self, env_ids): fixed_state = torch.cat( [ - wp.to_torch(self._fixed_asset.data.default_root_pose), - wp.to_torch(self._fixed_asset.data.default_root_vel), + self._fixed_asset.data.default_root_pose.torch, + self._fixed_asset.data.default_root_vel.torch, ], dim=-1, )[env_ids].clone() @@ -612,7 +612,7 @@ def _set_franka_to_default_pose(self, joints, env_ids): # gripper_width = self.cfg_task.held_asset_cfg.diameter / 2 * 1.25 # gripper_width = self.cfg_task.hand_width_max / 3.0 gripper_width = self.gripper_open_width - joint_pos = wp.to_torch(self._robot.data.default_joint_pos)[env_ids] + joint_pos = self._robot.data.default_joint_pos.torch[env_ids] joint_pos[:, 7:] = gripper_width # MIMIC joint_pos[:, :7] = torch.tensor(joints, device=self.device)[None, :] joint_vel = torch.zeros_like(joint_pos) @@ -637,8 +637,8 @@ def randomize_fixed_initial_state(self, env_ids): # (1.) Randomize fixed asset pose. fixed_state = torch.cat( [ - wp.to_torch(self._fixed_asset.data.default_root_pose), - wp.to_torch(self._fixed_asset.data.default_root_vel), + self._fixed_asset.data.default_root_pose.torch, + self._fixed_asset.data.default_root_vel.torch, ], dim=-1, )[env_ids].clone() @@ -678,7 +678,7 @@ def randomize_fixed_initial_state(self, env_ids): def randomize_held_initial_state(self, env_ids, pre_grasp): # Set plug pos to assembled state held_state = torch.cat( - [wp.to_torch(self._held_asset.data.default_root_pose), wp.to_torch(self._held_asset.data.default_root_vel)], + [self._held_asset.data.default_root_pose.torch, self._held_asset.data.default_root_vel.torch], dim=-1, ).clone() held_state[env_ids, 0:3] = self.fixed_pos[env_ids].clone() + self.scene.env_origins[env_ids] diff --git a/source/isaaclab_tasks/isaaclab_tasks/direct/cart_double_pendulum/cart_double_pendulum_env.py b/source/isaaclab_tasks/isaaclab_tasks/direct/cart_double_pendulum/cart_double_pendulum_env.py index cd3fca731952..f87ae20ce844 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/direct/cart_double_pendulum/cart_double_pendulum_env.py +++ b/source/isaaclab_tasks/isaaclab_tasks/direct/cart_double_pendulum/cart_double_pendulum_env.py @@ -9,7 +9,6 @@ from collections.abc import Sequence import torch -import warp as wp import isaaclab.sim as sim_utils from isaaclab.assets import Articulation, ArticulationCfg @@ -75,8 +74,8 @@ def __init__(self, cfg: CartDoublePendulumEnvCfg, render_mode: str | None = None self._pole_dof_idx, _ = self.robot.find_joints(self.cfg.pole_dof_name) self._pendulum_dof_idx, _ = self.robot.find_joints(self.cfg.pendulum_dof_name) - self.joint_pos = wp.to_torch(self.robot.data.joint_pos) - self.joint_vel = wp.to_torch(self.robot.data.joint_vel) + self.joint_pos = self.robot.data.joint_pos.torch + self.joint_vel = self.robot.data.joint_vel.torch def _setup_scene(self): self.robot = Articulation(self.cfg.robot_cfg) @@ -149,8 +148,8 @@ def _get_rewards(self) -> dict[str, torch.Tensor]: return total_reward def _get_dones(self) -> tuple[dict[str, torch.Tensor], dict[str, torch.Tensor]]: - self.joint_pos = wp.to_torch(self.robot.data.joint_pos) - self.joint_vel = wp.to_torch(self.robot.data.joint_vel) + self.joint_pos = self.robot.data.joint_pos.torch + self.joint_vel = self.robot.data.joint_vel.torch time_out = self.episode_length_buf >= self.max_episode_length - 1 out_of_bounds = torch.any(torch.abs(self.joint_pos[:, self._cart_dof_idx]) > self.cfg.max_cart_pos, dim=1) @@ -165,7 +164,7 @@ def _reset_idx(self, env_ids: Sequence[int] | None): env_ids = self.robot._ALL_INDICES super()._reset_idx(env_ids) - joint_pos = wp.to_torch(self.robot.data.default_joint_pos)[env_ids] + joint_pos = self.robot.data.default_joint_pos.torch[env_ids] joint_pos[:, self._pole_dof_idx] += sample_uniform( self.cfg.initial_pole_angle_range[0] * math.pi, self.cfg.initial_pole_angle_range[1] * math.pi, @@ -178,10 +177,10 @@ def _reset_idx(self, env_ids: Sequence[int] | None): joint_pos[:, self._pendulum_dof_idx].shape, joint_pos.device, ) - joint_vel = wp.to_torch(self.robot.data.default_joint_vel)[env_ids] + joint_vel = self.robot.data.default_joint_vel.torch[env_ids] - default_root_pose = wp.to_torch(self.robot.data.default_root_pose)[env_ids] - default_root_vel = wp.to_torch(self.robot.data.default_root_vel)[env_ids] + default_root_pose = self.robot.data.default_root_pose.torch[env_ids] + default_root_vel = self.robot.data.default_root_vel.torch[env_ids] default_root_pose[:, :3] += self.scene.env_origins[env_ids] self.joint_pos[env_ids] = joint_pos diff --git a/source/isaaclab_tasks/isaaclab_tasks/direct/cartpole/cartpole_camera_env.py b/source/isaaclab_tasks/isaaclab_tasks/direct/cartpole/cartpole_camera_env.py index 3c2a896766ea..18215b4788df 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/direct/cartpole/cartpole_camera_env.py +++ b/source/isaaclab_tasks/isaaclab_tasks/direct/cartpole/cartpole_camera_env.py @@ -10,7 +10,6 @@ from typing import TYPE_CHECKING import torch -import warp as wp import isaaclab.sim as sim_utils from isaaclab.assets import Articulation @@ -58,8 +57,8 @@ def __init__( self._pole_dof_idx, _ = self._cartpole.find_joints(self.cfg.pole_dof_name) self.action_scale = self.cfg.action_scale - self.joint_pos = wp.to_torch(self._cartpole.data.joint_pos) - self.joint_vel = wp.to_torch(self._cartpole.data.joint_vel) + self.joint_pos = self._cartpole.data.joint_pos.torch + self.joint_vel = self._cartpole.data.joint_vel.torch if len(self.cfg.tiled_camera.data_types) != 1: raise ValueError( @@ -141,8 +140,8 @@ def _get_rewards(self) -> torch.Tensor: return total_reward def _get_dones(self) -> tuple[torch.Tensor, torch.Tensor]: - self.joint_pos = wp.to_torch(self._cartpole.data.joint_pos) - self.joint_vel = wp.to_torch(self._cartpole.data.joint_vel) + self.joint_pos = self._cartpole.data.joint_pos.torch + self.joint_vel = self._cartpole.data.joint_vel.torch time_out = self.episode_length_buf >= self.max_episode_length - 1 out_of_bounds = torch.any(torch.abs(self.joint_pos[:, self._cart_dof_idx]) > self.cfg.max_cart_pos, dim=1) @@ -154,18 +153,18 @@ def _reset_idx(self, env_ids: Sequence[int] | None): env_ids = self._cartpole._ALL_INDICES super()._reset_idx(env_ids) - joint_pos = wp.to_torch(self._cartpole.data.default_joint_pos)[env_ids] + joint_pos = self._cartpole.data.default_joint_pos.torch[env_ids] joint_pos[:, self._pole_dof_idx] += sample_uniform( self.cfg.initial_pole_angle_range[0] * math.pi, self.cfg.initial_pole_angle_range[1] * math.pi, joint_pos[:, self._pole_dof_idx].shape, joint_pos.device, ) - joint_vel = wp.to_torch(self._cartpole.data.default_joint_vel)[env_ids] + joint_vel = self._cartpole.data.default_joint_vel.torch[env_ids] - default_root_pose = wp.to_torch(self._cartpole.data.default_root_pose)[env_ids] + default_root_pose = self._cartpole.data.default_root_pose.torch[env_ids] default_root_pose[:, :3] += self.scene.env_origins[env_ids] - default_root_vel = wp.to_torch(self._cartpole.data.default_root_vel)[env_ids] + default_root_vel = self._cartpole.data.default_root_vel.torch[env_ids] self.joint_pos[env_ids] = joint_pos self.joint_vel[env_ids] = joint_vel diff --git a/source/isaaclab_tasks/isaaclab_tasks/direct/cartpole/cartpole_env.py b/source/isaaclab_tasks/isaaclab_tasks/direct/cartpole/cartpole_env.py index 6c14fa8c2e18..70bfedf91795 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/direct/cartpole/cartpole_env.py +++ b/source/isaaclab_tasks/isaaclab_tasks/direct/cartpole/cartpole_env.py @@ -10,7 +10,6 @@ from typing import TYPE_CHECKING import torch -import warp as wp import isaaclab.sim as sim_utils from isaaclab.assets import Articulation @@ -32,8 +31,8 @@ def __init__(self, cfg: CartpoleEnvCfg, render_mode: str | None = None, **kwargs self._pole_dof_idx, _ = self.cartpole.find_joints(self.cfg.pole_dof_name) self.action_scale = self.cfg.action_scale - self.joint_pos = wp.to_torch(self.cartpole.data.joint_pos) - self.joint_vel = wp.to_torch(self.cartpole.data.joint_vel) + self.joint_pos = self.cartpole.data.joint_pos.torch + self.joint_vel = self.cartpole.data.joint_vel.torch def _setup_scene(self): self.cartpole = Articulation(self.cfg.robot_cfg) @@ -85,8 +84,8 @@ def _get_rewards(self) -> torch.Tensor: return total_reward def _get_dones(self) -> tuple[torch.Tensor, torch.Tensor]: - self.joint_pos = wp.to_torch(self.cartpole.data.joint_pos) - self.joint_vel = wp.to_torch(self.cartpole.data.joint_vel) + self.joint_pos = self.cartpole.data.joint_pos.torch + self.joint_vel = self.cartpole.data.joint_vel.torch time_out = self.episode_length_buf >= self.max_episode_length - 1 out_of_bounds = torch.any(torch.abs(self.joint_pos[:, self._cart_dof_idx]) > self.cfg.max_cart_pos, dim=1) @@ -98,18 +97,18 @@ def _reset_idx(self, env_ids: Sequence[int] | None): env_ids = self.cartpole._ALL_INDICES super()._reset_idx(env_ids) - joint_pos = wp.to_torch(self.cartpole.data.default_joint_pos)[env_ids].clone() + joint_pos = self.cartpole.data.default_joint_pos.torch[env_ids].clone() joint_pos[:, self._pole_dof_idx] += sample_uniform( self.cfg.initial_pole_angle_range[0] * math.pi, self.cfg.initial_pole_angle_range[1] * math.pi, joint_pos[:, self._pole_dof_idx].shape, joint_pos.device, ) - joint_vel = wp.to_torch(self.cartpole.data.default_joint_vel)[env_ids].clone() + joint_vel = self.cartpole.data.default_joint_vel.torch[env_ids].clone() - default_root_pose = wp.to_torch(self.cartpole.data.default_root_pose)[env_ids].clone() + default_root_pose = self.cartpole.data.default_root_pose.torch[env_ids].clone() default_root_pose[:, :3] += self.scene.env_origins[env_ids] - default_root_vel = wp.to_torch(self.cartpole.data.default_root_vel)[env_ids].clone() + default_root_vel = self.cartpole.data.default_root_vel.torch[env_ids].clone() self.joint_pos[env_ids] = joint_pos self.joint_vel[env_ids] = joint_vel diff --git a/source/isaaclab_tasks/isaaclab_tasks/direct/factory/factory_env.py b/source/isaaclab_tasks/isaaclab_tasks/direct/factory/factory_env.py index 10c1ad3e1a1a..6cd11408dfae 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/direct/factory/factory_env.py +++ b/source/isaaclab_tasks/isaaclab_tasks/direct/factory/factory_env.py @@ -118,18 +118,18 @@ def _setup_scene(self): def _compute_intermediate_values(self, dt): """Get values computed from raw tensors. This includes adding noise.""" # TODO: A lot of these can probably only be set once? - self.fixed_pos = wp.to_torch(self._fixed_asset.data.root_pos_w) - self.scene.env_origins - self.fixed_quat = wp.to_torch(self._fixed_asset.data.root_quat_w) + self.fixed_pos = self._fixed_asset.data.root_pos_w.torch - self.scene.env_origins + self.fixed_quat = self._fixed_asset.data.root_quat_w.torch - self.held_pos = wp.to_torch(self._held_asset.data.root_pos_w) - self.scene.env_origins - self.held_quat = wp.to_torch(self._held_asset.data.root_quat_w) + self.held_pos = self._held_asset.data.root_pos_w.torch - self.scene.env_origins + self.held_quat = self._held_asset.data.root_quat_w.torch self.fingertip_midpoint_pos = ( - wp.to_torch(self._robot.data.body_pos_w)[:, self.fingertip_body_idx] - self.scene.env_origins + self._robot.data.body_pos_w.torch[:, self.fingertip_body_idx] - self.scene.env_origins ) - self.fingertip_midpoint_quat = wp.to_torch(self._robot.data.body_quat_w)[:, self.fingertip_body_idx] - self.fingertip_midpoint_linvel = wp.to_torch(self._robot.data.body_lin_vel_w)[:, self.fingertip_body_idx] - self.fingertip_midpoint_angvel = wp.to_torch(self._robot.data.body_ang_vel_w)[:, self.fingertip_body_idx] + self.fingertip_midpoint_quat = self._robot.data.body_quat_w.torch[:, self.fingertip_body_idx] + self.fingertip_midpoint_linvel = self._robot.data.body_lin_vel_w.torch[:, self.fingertip_body_idx] + self.fingertip_midpoint_angvel = self._robot.data.body_ang_vel_w.torch[:, self.fingertip_body_idx] jacobians = wp.to_torch(self._robot.root_view.get_jacobians()) @@ -137,8 +137,8 @@ def _compute_intermediate_values(self, dt): self.right_finger_jacobian = jacobians[:, self.right_finger_body_idx - 1, 0:6, 0:7] self.fingertip_midpoint_jacobian = (self.left_finger_jacobian + self.right_finger_jacobian) * 0.5 self.arm_mass_matrix = wp.to_torch(self._robot.root_view.get_generalized_mass_matrices())[:, 0:7, 0:7] - self.joint_pos = wp.to_torch(self._robot.data.joint_pos).clone() - self.joint_vel = wp.to_torch(self._robot.data.joint_vel).clone() + self.joint_pos = self._robot.data.joint_pos.torch.clone() + self.joint_vel = self._robot.data.joint_vel.torch.clone() # Finite-differencing results in more reliable velocity estimates. self.ee_linvel_fd = (self.fingertip_midpoint_pos - self.prev_fingertip_pos) / dt @@ -499,16 +499,16 @@ def _reset_idx(self, env_ids): def _set_assets_to_default_pose(self, env_ids): """Move assets to default pose before randomization.""" - held_pose = wp.to_torch(self._held_asset.data.default_root_pose).clone()[env_ids] - held_vel = wp.to_torch(self._held_asset.data.default_root_vel).clone()[env_ids] + held_pose = self._held_asset.data.default_root_pose.torch.clone()[env_ids] + held_vel = self._held_asset.data.default_root_vel.torch.clone()[env_ids] held_pose[:, 0:3] += self.scene.env_origins[env_ids] held_vel[:] = 0.0 self._held_asset.write_root_pose_to_sim_index(root_pose=held_pose, env_ids=env_ids) self._held_asset.write_root_velocity_to_sim_index(root_velocity=held_vel, env_ids=env_ids) self._held_asset.reset() - fixed_pose = wp.to_torch(self._fixed_asset.data.default_root_pose).clone()[env_ids] - fixed_vel = wp.to_torch(self._fixed_asset.data.default_root_vel).clone()[env_ids] + fixed_pose = self._fixed_asset.data.default_root_pose.torch.clone()[env_ids] + fixed_vel = self._fixed_asset.data.default_root_vel.torch.clone()[env_ids] fixed_pose[:, 0:3] += self.scene.env_origins[env_ids] fixed_vel[:] = 0.0 self._fixed_asset.write_root_pose_to_sim_index(root_pose=fixed_pose, env_ids=env_ids) @@ -592,7 +592,7 @@ def get_handheld_asset_relative_pose(self): def _set_franka_to_default_pose(self, joints, env_ids): """Return Franka to its default joint position.""" gripper_width = self.cfg_task.held_asset_cfg.diameter / 2 * 1.25 - joint_pos = wp.to_torch(self._robot.data.default_joint_pos)[env_ids] + joint_pos = self._robot.data.default_joint_pos.torch[env_ids] joint_pos[:, 7:] = gripper_width # MIMIC joint_pos[:, :7] = torch.tensor(joints, device=self.device)[None, :] joint_vel = torch.zeros_like(joint_pos) @@ -624,8 +624,8 @@ def randomize_initial_state(self, env_ids): physics_sim_view.set_gravity(carb.Float3(0.0, 0.0, 0.0)) # (1.) Randomize fixed asset pose. - fixed_pose = wp.to_torch(self._fixed_asset.data.default_root_pose).clone()[env_ids] - fixed_vel = wp.to_torch(self._fixed_asset.data.default_root_vel).clone()[env_ids] + fixed_pose = self._fixed_asset.data.default_root_pose.torch.clone()[env_ids] + fixed_vel = self._fixed_asset.data.default_root_vel.torch.clone()[env_ids] # (1.a.) Position rand_sample = torch.rand((len(env_ids), 3), dtype=torch.float32, device=self.device) fixed_pos_init_rand = 2 * (rand_sample - 0.5) # [-1, 1] @@ -732,16 +732,16 @@ def randomize_initial_state(self, env_ids): # Add flanking gears after servo (so arm doesn't move them). if self.cfg_task.name == "gear_mesh" and self.cfg_task.add_flanking_gears: - small_gear_pose = wp.to_torch(self._small_gear_asset.data.default_root_pose).clone()[env_ids] - small_gear_vel = wp.to_torch(self._small_gear_asset.data.default_root_vel).clone()[env_ids] + small_gear_pose = self._small_gear_asset.data.default_root_pose.torch.clone()[env_ids] + small_gear_vel = self._small_gear_asset.data.default_root_vel.torch.clone()[env_ids] small_gear_pose[:, 0:7] = fixed_pose[:, 0:7] small_gear_vel[:] = 0.0 # vel self._small_gear_asset.write_root_pose_to_sim_index(root_pose=small_gear_pose, env_ids=env_ids) self._small_gear_asset.write_root_velocity_to_sim_index(root_velocity=small_gear_vel, env_ids=env_ids) self._small_gear_asset.reset() - large_gear_pose = wp.to_torch(self._large_gear_asset.data.default_root_pose).clone()[env_ids] - large_gear_vel = wp.to_torch(self._large_gear_asset.data.default_root_vel).clone()[env_ids] + large_gear_pose = self._large_gear_asset.data.default_root_pose.torch.clone()[env_ids] + large_gear_vel = self._large_gear_asset.data.default_root_vel.torch.clone()[env_ids] large_gear_pose[:, 0:7] = fixed_pose[:, 0:7] large_gear_vel[:] = 0.0 # vel self._large_gear_asset.write_root_pose_to_sim_index(root_pose=large_gear_pose, env_ids=env_ids) @@ -783,8 +783,8 @@ def randomize_initial_state(self, env_ids): torch.tensor([0.0, 0.0, 0.0, 1.0], device=self.device).unsqueeze(0).repeat(self.num_envs, 1), ) - held_pose = wp.to_torch(self._held_asset.data.default_root_pose).clone() - held_vel = wp.to_torch(self._held_asset.data.default_root_vel).clone() + held_pose = self._held_asset.data.default_root_pose.torch.clone() + held_vel = self._held_asset.data.default_root_vel.torch.clone() held_pose[:, 0:3] = translated_held_asset_pos + self.scene.env_origins held_pose[:, 3:7] = translated_held_asset_quat held_vel[:] = 0.0 diff --git a/source/isaaclab_tasks/isaaclab_tasks/direct/franka_cabinet/franka_cabinet_env.py b/source/isaaclab_tasks/isaaclab_tasks/direct/franka_cabinet/franka_cabinet_env.py index 0106e66c5950..dda1de146727 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/direct/franka_cabinet/franka_cabinet_env.py +++ b/source/isaaclab_tasks/isaaclab_tasks/direct/franka_cabinet/franka_cabinet_env.py @@ -54,12 +54,8 @@ def get_env_local_pose(env_pos: torch.Tensor, xformable: UsdGeom.Xformable, devi self.dt = self.cfg.sim.dt * self.cfg.decimation # create auxiliary variables for computing applied action, observations and rewards - self.robot_dof_lower_limits = wp.to_torch(self._robot.data.soft_joint_pos_limits)[0, :, 0].to( - device=self.device - ) - self.robot_dof_upper_limits = wp.to_torch(self._robot.data.soft_joint_pos_limits)[0, :, 1].to( - device=self.device - ) + self.robot_dof_lower_limits = self._robot.data.soft_joint_pos_limits.torch[0, :, 0].to(device=self.device) + self.robot_dof_upper_limits = self._robot.data.soft_joint_pos_limits.torch[0, :, 1].to(device=self.device) self.robot_dof_speed_scales = torch.ones_like(self.robot_dof_lower_limits) self.robot_dof_speed_scales[self._robot.find_joints("panda_finger_joint1")[0]] = 0.1 @@ -159,19 +155,19 @@ def _apply_action(self): # post-physics step calls def _get_dones(self) -> tuple[torch.Tensor, torch.Tensor]: - terminated = wp.to_torch(self._cabinet.data.joint_pos)[:, self.drawer_joint_idx] > 0.39 + terminated = self._cabinet.data.joint_pos.torch[:, self.drawer_joint_idx] > 0.39 truncated = self.episode_length_buf >= self.max_episode_length - 1 return terminated, truncated def _get_rewards(self) -> torch.Tensor: # Refresh the intermediate values after the physics steps self._compute_intermediate_values() - robot_left_finger_pos = wp.to_torch(self._robot.data.body_pos_w)[:, self.left_finger_link_idx] - robot_right_finger_pos = wp.to_torch(self._robot.data.body_pos_w)[:, self.right_finger_link_idx] + robot_left_finger_pos = self._robot.data.body_pos_w.torch[:, self.left_finger_link_idx] + robot_right_finger_pos = self._robot.data.body_pos_w.torch[:, self.right_finger_link_idx] return self._compute_rewards( self.actions, - wp.to_torch(self._cabinet.data.joint_pos), + self._cabinet.data.joint_pos.torch, self.robot_grasp_pos, self.drawer_grasp_pos, self.robot_grasp_rot, @@ -188,13 +184,13 @@ def _get_rewards(self) -> torch.Tensor: self.cfg.open_reward_scale, self.cfg.action_penalty_scale, self.cfg.finger_reward_scale, - wp.to_torch(self._robot.data.joint_pos), + self._robot.data.joint_pos.torch, ) def _reset_idx(self, env_ids: torch.Tensor | None): super()._reset_idx(env_ids) # robot state - joint_pos = wp.to_torch(self._robot.data.default_joint_pos)[env_ids] + sample_uniform( + joint_pos = self._robot.data.default_joint_pos.torch[env_ids] + sample_uniform( -0.125, 0.125, (len(env_ids), self._robot.num_joints), @@ -217,7 +213,7 @@ def _reset_idx(self, env_ids: torch.Tensor | None): def _get_observations(self) -> dict: dof_pos_scaled = ( 2.0 - * (wp.to_torch(self._robot.data.joint_pos) - self.robot_dof_lower_limits) + * (self._robot.data.joint_pos.torch - self.robot_dof_lower_limits) / (self.robot_dof_upper_limits - self.robot_dof_lower_limits) - 1.0 ) @@ -226,10 +222,10 @@ def _get_observations(self) -> dict: obs = torch.cat( ( dof_pos_scaled, - wp.to_torch(self._robot.data.joint_vel) * self.cfg.dof_velocity_scale, + self._robot.data.joint_vel.torch * self.cfg.dof_velocity_scale, to_target, - wp.to_torch(self._cabinet.data.joint_pos)[:, self.drawer_joint_idx].unsqueeze(-1), - wp.to_torch(self._cabinet.data.joint_vel)[:, self.drawer_joint_idx].unsqueeze(-1), + self._cabinet.data.joint_pos.torch[:, self.drawer_joint_idx].unsqueeze(-1), + self._cabinet.data.joint_vel.torch[:, self.drawer_joint_idx].unsqueeze(-1), ), dim=-1, ) @@ -241,10 +237,10 @@ def _compute_intermediate_values(self, env_ids: torch.Tensor | None = None): if env_ids is None: env_ids = wp.to_torch(self._robot._ALL_INDICES) - hand_pos = wp.to_torch(self._robot.data.body_pos_w)[env_ids, self.hand_link_idx] - hand_rot = wp.to_torch(self._robot.data.body_quat_w)[env_ids, self.hand_link_idx] - drawer_pos = wp.to_torch(self._cabinet.data.body_pos_w)[env_ids, self.drawer_link_idx] - drawer_rot = wp.to_torch(self._cabinet.data.body_quat_w)[env_ids, self.drawer_link_idx] + hand_pos = self._robot.data.body_pos_w.torch[env_ids, self.hand_link_idx] + hand_rot = self._robot.data.body_quat_w.torch[env_ids, self.hand_link_idx] + drawer_pos = self._cabinet.data.body_pos_w.torch[env_ids, self.drawer_link_idx] + drawer_rot = self._cabinet.data.body_quat_w.torch[env_ids, self.drawer_link_idx] ( self.robot_grasp_rot[env_ids], self.robot_grasp_pos[env_ids], diff --git a/source/isaaclab_tasks/isaaclab_tasks/direct/humanoid_amp/humanoid_amp_env.py b/source/isaaclab_tasks/isaaclab_tasks/direct/humanoid_amp/humanoid_amp_env.py index ccd82d79dc5c..7c0fe5036132 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/direct/humanoid_amp/humanoid_amp_env.py +++ b/source/isaaclab_tasks/isaaclab_tasks/direct/humanoid_amp/humanoid_amp_env.py @@ -27,7 +27,7 @@ def __init__(self, cfg: HumanoidAmpEnvCfg, render_mode: str | None = None, **kwa super().__init__(cfg, render_mode, **kwargs) # action offset and scale - soft_joint_pos_limits = wp.to_torch(self.robot.data.soft_joint_pos_limits) + soft_joint_pos_limits = self.robot.data.soft_joint_pos_limits.torch dof_lower_limits = soft_joint_pos_limits[0, :, 0] dof_upper_limits = soft_joint_pos_limits[0, :, 1] self.action_offset = 0.5 * (dof_upper_limits + dof_lower_limits) @@ -86,13 +86,13 @@ def _apply_action(self): def _get_observations(self) -> dict: # build task observation obs = compute_obs( - wp.to_torch(self.robot.data.joint_pos), - wp.to_torch(self.robot.data.joint_vel), - wp.to_torch(self.robot.data.body_pos_w)[:, self.ref_body_index], - wp.to_torch(self.robot.data.body_quat_w)[:, self.ref_body_index], - wp.to_torch(self.robot.data.body_lin_vel_w)[:, self.ref_body_index], - wp.to_torch(self.robot.data.body_ang_vel_w)[:, self.ref_body_index], - wp.to_torch(self.robot.data.body_pos_w)[:, self.key_body_indexes], + self.robot.data.joint_pos.torch, + self.robot.data.joint_vel.torch, + self.robot.data.body_pos_w.torch[:, self.ref_body_index], + self.robot.data.body_quat_w.torch[:, self.ref_body_index], + self.robot.data.body_lin_vel_w.torch[:, self.ref_body_index], + self.robot.data.body_ang_vel_w.torch[:, self.ref_body_index], + self.robot.data.body_pos_w.torch[:, self.key_body_indexes], ) # update AMP observation history @@ -110,7 +110,7 @@ def _get_rewards(self) -> torch.Tensor: def _get_dones(self) -> tuple[torch.Tensor, torch.Tensor]: time_out = self.episode_length_buf >= self.max_episode_length - 1 if self.cfg.early_termination: - died = wp.to_torch(self.robot.data.body_pos_w)[:, self.ref_body_index, 2] < self.cfg.termination_height + died = self.robot.data.body_pos_w.torch[:, self.ref_body_index, 2] < self.cfg.termination_height else: died = torch.zeros_like(time_out) return died, time_out @@ -142,12 +142,12 @@ def _reset_idx(self, env_ids: torch.Tensor | None): # reset strategies def _reset_strategy_default(self, env_ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - default_root_pose = wp.to_torch(self.robot.data.default_root_pose)[env_ids].clone() - default_root_vel = wp.to_torch(self.robot.data.default_root_vel)[env_ids].clone() + default_root_pose = self.robot.data.default_root_pose.torch[env_ids].clone() + default_root_vel = self.robot.data.default_root_vel.torch[env_ids].clone() default_root_pose[:, :3] += self.scene.env_origins[env_ids] root_state = torch.cat([default_root_pose, default_root_vel], dim=-1) - joint_pos = self.robot.data.default_joint_pos[env_ids].clone() - joint_vel = self.robot.data.default_joint_vel[env_ids].clone() + joint_pos = self.robot.data.default_joint_pos.torch[env_ids].clone() + joint_vel = self.robot.data.default_joint_vel.torch[env_ids].clone() return root_state, joint_pos, joint_vel def _reset_strategy_random( @@ -170,8 +170,8 @@ def _reset_strategy_random( motion_torso_index = self._motion_loader.get_body_index(["torso"])[0] root_state = torch.cat( [ - wp.to_torch(self.robot.data.default_root_pose)[env_ids], - wp.to_torch(self.robot.data.default_root_vel)[env_ids], + self.robot.data.default_root_pose.torch[env_ids], + self.robot.data.default_root_vel.torch[env_ids], ], dim=-1, ).clone() diff --git a/source/isaaclab_tasks/isaaclab_tasks/direct/inhand_manipulation/inhand_manipulation_env.py b/source/isaaclab_tasks/isaaclab_tasks/direct/inhand_manipulation/inhand_manipulation_env.py index 66a09e7e3cde..adbf96b97d80 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/direct/inhand_manipulation/inhand_manipulation_env.py +++ b/source/isaaclab_tasks/isaaclab_tasks/direct/inhand_manipulation/inhand_manipulation_env.py @@ -11,7 +11,6 @@ import numpy as np import torch -import warp as wp import isaaclab.sim as sim_utils from isaaclab.assets import Articulation, RigidObject @@ -52,14 +51,14 @@ def __init__(self, cfg: AllegroHandEnvCfg | ShadowHandEnvCfg, render_mode: str | self.num_fingertips = len(self.finger_bodies) # joint limits - joint_pos_limits = wp.to_torch(self.hand.data.joint_limits).to(self.device) + joint_pos_limits = self.hand.data.joint_limits.torch.to(self.device) self.hand_dof_lower_limits = joint_pos_limits[..., 0] self.hand_dof_upper_limits = joint_pos_limits[..., 1] # track goal resets self.reset_goal_buf = torch.zeros(self.num_envs, dtype=torch.bool, device=self.device) # used to compare object position - self.in_hand_pos = wp.to_torch(self.object.data.default_root_pose)[:, 0:3].clone() + self.in_hand_pos = self.object.data.default_root_pose.torch[:, 0:3].clone() self.in_hand_pos[:, 2] -= 0.04 # default goal positions self.goal_rot = torch.zeros((self.num_envs, 4), dtype=torch.float, device=self.device) @@ -137,9 +136,7 @@ def _get_observations(self) -> dict: if self.cfg.asymmetric_obs: # Newton does not implement body_incoming_joint_wrench_b; fall back to zeros. try: - self.fingertip_force_sensors = wp.to_torch(self.hand.data.body_incoming_joint_wrench_b)[ - :, self.finger_bodies - ] + self.fingertip_force_sensors = self.hand.data.body_incoming_joint_wrench_b.torch[:, self.finger_bodies] except NotImplementedError: self.fingertip_force_sensors = torch.zeros( self.num_envs, len(self.finger_bodies), 6, dtype=torch.float32, device=self.device @@ -228,8 +225,8 @@ def _reset_idx(self, env_ids: Sequence[int]): self._reset_target_pose(env_ids) # reset object - object_default_pose = wp.to_torch(self.object.data.default_root_pose).clone()[env_ids] - object_default_vel = wp.to_torch(self.object.data.default_root_vel).clone()[env_ids] + object_default_pose = self.object.data.default_root_pose.torch.clone()[env_ids] + object_default_vel = self.object.data.default_root_vel.torch.clone()[env_ids] pos_noise = sample_uniform(-1.0, 1.0, (len(env_ids), 3), device=self.device) # global object positions object_default_pose[:, 0:3] = ( @@ -246,15 +243,15 @@ def _reset_idx(self, env_ids: Sequence[int]): self._write_obj_root_vel(root_velocity=object_default_vel, env_ids=env_ids) # reset hand - delta_max = self.hand_dof_upper_limits[env_ids] - wp.to_torch(self.hand.data.default_joint_pos)[env_ids] - delta_min = self.hand_dof_lower_limits[env_ids] - wp.to_torch(self.hand.data.default_joint_pos)[env_ids] + delta_max = self.hand_dof_upper_limits[env_ids] - self.hand.data.default_joint_pos.torch[env_ids] + delta_min = self.hand_dof_lower_limits[env_ids] - self.hand.data.default_joint_pos.torch[env_ids] dof_pos_noise = sample_uniform(-1.0, 1.0, (len(env_ids), self.num_hand_dofs), device=self.device) rand_delta = delta_min + (delta_max - delta_min) * 0.5 * dof_pos_noise - dof_pos = wp.to_torch(self.hand.data.default_joint_pos)[env_ids] + self.cfg.reset_dof_pos_noise * rand_delta + dof_pos = self.hand.data.default_joint_pos.torch[env_ids] + self.cfg.reset_dof_pos_noise * rand_delta dof_vel_noise = sample_uniform(-1.0, 1.0, (len(env_ids), self.num_hand_dofs), device=self.device) - dof_vel = wp.to_torch(self.hand.data.default_joint_vel)[env_ids] + self.cfg.reset_dof_vel_noise * dof_vel_noise + dof_vel = self.hand.data.default_joint_vel.torch[env_ids] + self.cfg.reset_dof_vel_noise * dof_vel_noise self.prev_targets[env_ids] = dof_pos self.cur_targets[env_ids] = dof_pos @@ -283,22 +280,22 @@ def _reset_target_pose(self, env_ids): def _compute_intermediate_values(self): # data for hand - self.fingertip_pos = wp.to_torch(self.hand.data.body_pos_w)[:, self.finger_bodies] - self.fingertip_rot = wp.to_torch(self.hand.data.body_quat_w)[:, self.finger_bodies] + self.fingertip_pos = self.hand.data.body_pos_w.torch[:, self.finger_bodies] + self.fingertip_rot = self.hand.data.body_quat_w.torch[:, self.finger_bodies] self.fingertip_pos -= self.scene.env_origins.repeat((1, self.num_fingertips)).reshape( self.num_envs, self.num_fingertips, 3 ) - self.fingertip_velocities = wp.to_torch(self.hand.data.body_vel_w)[:, self.finger_bodies] + self.fingertip_velocities = self.hand.data.body_vel_w.torch[:, self.finger_bodies] - self.hand_dof_pos = wp.to_torch(self.hand.data.joint_pos) - self.hand_dof_vel = wp.to_torch(self.hand.data.joint_vel) + self.hand_dof_pos = self.hand.data.joint_pos.torch + self.hand_dof_vel = self.hand.data.joint_vel.torch # data for object - self.object_pos = wp.to_torch(self.object.data.root_pos_w) - self.scene.env_origins - self.object_rot = wp.to_torch(self.object.data.root_quat_w) - self.object_velocities = wp.to_torch(self.object.data.root_vel_w) - self.object_linvel = wp.to_torch(self.object.data.root_lin_vel_w) - self.object_angvel = wp.to_torch(self.object.data.root_ang_vel_w) + self.object_pos = self.object.data.root_pos_w.torch - self.scene.env_origins + self.object_rot = self.object.data.root_quat_w.torch + self.object_velocities = self.object.data.root_vel_w.torch + self.object_linvel = self.object.data.root_lin_vel_w.torch + self.object_angvel = self.object.data.root_ang_vel_w.torch def compute_reduced_observations(self): # Per https://arxiv.org/pdf/1808.00177.pdf Table 2 diff --git a/source/isaaclab_tasks/isaaclab_tasks/direct/locomotion/locomotion_env.py b/source/isaaclab_tasks/isaaclab_tasks/direct/locomotion/locomotion_env.py index 87d7a6d8b3d8..fc5356ad17cf 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/direct/locomotion/locomotion_env.py +++ b/source/isaaclab_tasks/isaaclab_tasks/direct/locomotion/locomotion_env.py @@ -133,14 +133,14 @@ def _apply_action(self): def _compute_intermediate_values(self): self.torso_position, self.torso_rotation = ( - wp.to_torch(self.robot.data.root_pos_w), - wp.to_torch(self.robot.data.root_quat_w), + self.robot.data.root_pos_w.torch, + self.robot.data.root_quat_w.torch, ) self.velocity, self.ang_velocity = ( - wp.to_torch(self.robot.data.root_lin_vel_w), - wp.to_torch(self.robot.data.root_ang_vel_w), + self.robot.data.root_lin_vel_w.torch, + self.robot.data.root_ang_vel_w.torch, ) - self.dof_pos, self.dof_vel = wp.to_torch(self.robot.data.joint_pos), wp.to_torch(self.robot.data.joint_vel) + self.dof_pos, self.dof_vel = self.robot.data.joint_pos.torch, self.robot.data.joint_vel.torch ( self.up_proj, @@ -163,8 +163,8 @@ def _compute_intermediate_values(self): self.velocity, self.ang_velocity, self.dof_pos, - wp.to_torch(self.robot.data.soft_joint_pos_limits)[0, :, 0], - wp.to_torch(self.robot.data.soft_joint_pos_limits)[0, :, 1], + self.robot.data.soft_joint_pos_limits.torch[0, :, 0], + self.robot.data.soft_joint_pos_limits.torch[0, :, 1], self.inv_start_rot, self.basis_vec0, self.basis_vec1, @@ -226,10 +226,10 @@ def _reset_idx(self, env_ids: torch.Tensor | None): self.robot.reset(env_ids) super()._reset_idx(env_ids) - joint_pos = wp.to_torch(self.robot.data.default_joint_pos)[env_ids].clone() - joint_vel = wp.to_torch(self.robot.data.default_joint_vel)[env_ids].clone() - default_root_pose = wp.to_torch(self.robot.data.default_root_pose)[env_ids].clone() - default_root_vel = wp.to_torch(self.robot.data.default_root_vel)[env_ids].clone() + joint_pos = self.robot.data.default_joint_pos.torch[env_ids].clone() + joint_vel = self.robot.data.default_joint_vel.torch[env_ids].clone() + default_root_pose = self.robot.data.default_root_pose.torch[env_ids].clone() + default_root_vel = self.robot.data.default_root_vel.torch[env_ids].clone() default_root_pose[:, :3] += self.scene.env_origins[env_ids] self.robot.write_root_pose_to_sim_index(root_pose=default_root_pose, env_ids=env_ids) diff --git a/source/isaaclab_tasks/isaaclab_tasks/direct/quadcopter/quadcopter_env.py b/source/isaaclab_tasks/isaaclab_tasks/direct/quadcopter/quadcopter_env.py index 51ad02d88958..1fc0fa8f6e48 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/direct/quadcopter/quadcopter_env.py +++ b/source/isaaclab_tasks/isaaclab_tasks/direct/quadcopter/quadcopter_env.py @@ -100,13 +100,13 @@ def _apply_action(self): def _get_observations(self) -> dict: desired_pos_b, _ = subtract_frame_transforms( - wp.to_torch(self._robot.data.root_pos_w), wp.to_torch(self._robot.data.root_quat_w), self._desired_pos_w + self._robot.data.root_pos_w.torch, self._robot.data.root_quat_w.torch, self._desired_pos_w ) obs = torch.cat( [ - wp.to_torch(self._robot.data.root_lin_vel_b), - wp.to_torch(self._robot.data.root_ang_vel_b), - wp.to_torch(self._robot.data.projected_gravity_b), + self._robot.data.root_lin_vel_b.torch, + self._robot.data.root_ang_vel_b.torch, + self._robot.data.projected_gravity_b.torch, desired_pos_b, ], dim=-1, @@ -115,9 +115,9 @@ def _get_observations(self) -> dict: return observations def _get_rewards(self) -> torch.Tensor: - lin_vel = torch.sum(torch.square(wp.to_torch(self._robot.data.root_lin_vel_b)), dim=1) - ang_vel = torch.sum(torch.square(wp.to_torch(self._robot.data.root_ang_vel_b)), dim=1) - distance_to_goal = torch.linalg.norm(self._desired_pos_w - wp.to_torch(self._robot.data.root_pos_w), dim=1) + lin_vel = torch.sum(torch.square(self._robot.data.root_lin_vel_b.torch), dim=1) + ang_vel = torch.sum(torch.square(self._robot.data.root_ang_vel_b.torch), dim=1) + distance_to_goal = torch.linalg.norm(self._desired_pos_w - self._robot.data.root_pos_w.torch, dim=1) distance_to_goal_mapped = 1 - torch.tanh(distance_to_goal / 0.8) rewards = { "lin_vel": lin_vel * self.cfg.lin_vel_reward_scale * self.step_dt, @@ -133,7 +133,7 @@ def _get_rewards(self) -> torch.Tensor: def _get_dones(self) -> tuple[torch.Tensor, torch.Tensor]: time_out = self.episode_length_buf >= self.max_episode_length - 1 died = torch.logical_or( - wp.to_torch(self._robot.data.root_pos_w)[:, 2] < 0.1, wp.to_torch(self._robot.data.root_pos_w)[:, 2] > 2.0 + self._robot.data.root_pos_w.torch[:, 2] < 0.1, self._robot.data.root_pos_w.torch[:, 2] > 2.0 ) return died, time_out @@ -143,7 +143,7 @@ def _reset_idx(self, env_ids: torch.Tensor | None): # Logging final_distance_to_goal = torch.linalg.norm( - self._desired_pos_w[env_ids] - wp.to_torch(self._robot.data.root_pos_w)[env_ids], dim=1 + self._desired_pos_w[env_ids] - self._robot.data.root_pos_w.torch[env_ids], dim=1 ).mean() extras = dict() for key in self._episode_sums.keys(): @@ -170,10 +170,10 @@ def _reset_idx(self, env_ids: torch.Tensor | None): self._desired_pos_w[env_ids, :2] += self._terrain.env_origins[env_ids, :2] self._desired_pos_w[env_ids, 2] = torch.zeros_like(self._desired_pos_w[env_ids, 2]).uniform_(0.5, 1.5) # Reset robot state - joint_pos = wp.to_torch(self._robot.data.default_joint_pos)[env_ids] - joint_vel = wp.to_torch(self._robot.data.default_joint_vel)[env_ids] - default_root_pose = wp.to_torch(self._robot.data.default_root_pose)[env_ids] - default_root_vel = wp.to_torch(self._robot.data.default_root_vel)[env_ids] + joint_pos = self._robot.data.default_joint_pos.torch[env_ids] + joint_vel = self._robot.data.default_joint_vel.torch[env_ids] + default_root_pose = self._robot.data.default_root_pose.torch[env_ids] + default_root_vel = self._robot.data.default_root_vel.torch[env_ids] default_root_pose[:, :3] += self._terrain.env_origins[env_ids] self._robot.write_root_pose_to_sim_index(root_pose=default_root_pose, env_ids=env_ids) self._robot.write_root_velocity_to_sim_index(root_velocity=default_root_vel, env_ids=env_ids) diff --git a/source/isaaclab_tasks/isaaclab_tasks/direct/shadow_hand/shadow_hand_vision_env.py b/source/isaaclab_tasks/isaaclab_tasks/direct/shadow_hand/shadow_hand_vision_env.py index 56f8c8fb8a56..b270f044ebf3 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/direct/shadow_hand/shadow_hand_vision_env.py +++ b/source/isaaclab_tasks/isaaclab_tasks/direct/shadow_hand/shadow_hand_vision_env.py @@ -9,7 +9,6 @@ from typing import TYPE_CHECKING import torch -import warp as wp import isaaclab.sim as sim_utils from isaaclab.assets import Articulation, RigidObject @@ -130,9 +129,7 @@ def _get_observations(self) -> dict: obs = torch.cat((state_obs, image_obs), dim=-1) # asymmetric critic states — Newton does not implement body_incoming_joint_wrench_b try: - self.fingertip_force_sensors = wp.to_torch(self.hand.data.body_incoming_joint_wrench_b)[ - :, self.finger_bodies - ] + self.fingertip_force_sensors = self.hand.data.body_incoming_joint_wrench_b.torch[:, self.finger_bodies] except NotImplementedError: self.fingertip_force_sensors = torch.zeros( self.num_envs, len(self.finger_bodies), 6, dtype=torch.float32, device=self.device diff --git a/source/isaaclab_tasks/isaaclab_tasks/direct/shadow_hand_over/shadow_hand_over_env.py b/source/isaaclab_tasks/isaaclab_tasks/direct/shadow_hand_over/shadow_hand_over_env.py index 7662fe728341..8b2598823d80 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/direct/shadow_hand_over/shadow_hand_over_env.py +++ b/source/isaaclab_tasks/isaaclab_tasks/direct/shadow_hand_over/shadow_hand_over_env.py @@ -69,7 +69,7 @@ def __init__(self, cfg: ShadowHandOverEnvCfg, render_mode: str | None = None, ** self.hand_dof_upper_limits = joint_pos_limits[..., 1] # used to compare object position - self.in_hand_pos = wp.to_torch(self.object.data.default_root_pose)[:, 0:3].clone() + self.in_hand_pos = self.object.data.default_root_pose.torch[:, 0:3].clone() self.in_hand_pos[:, 2] -= 0.04 # default goal positions self.goal_rot = torch.zeros((self.num_envs, 4), dtype=torch.float, device=self.device) @@ -311,8 +311,8 @@ def _reset_idx(self, env_ids: Sequence[int] | torch.Tensor | None): self._reset_target_pose(env_ids) # reset object - object_default_pose = wp.to_torch(self.object.data.default_root_pose).clone()[env_ids] - object_default_vel = wp.to_torch(self.object.data.default_root_vel).clone()[env_ids] + object_default_pose = self.object.data.default_root_pose.torch.clone()[env_ids] + object_default_vel = self.object.data.default_root_vel.torch.clone()[env_ids] pos_noise = sample_uniform(-1.0, 1.0, (len(env_ids), 3), device=self.device) object_default_pose[:, 0:3] = ( @@ -329,19 +329,15 @@ def _reset_idx(self, env_ids: Sequence[int] | torch.Tensor | None): self.object.write_root_velocity_to_sim_index(root_velocity=object_default_vel, env_ids=env_ids) # reset right hand - delta_max = self.hand_dof_upper_limits[env_ids] - wp.to_torch(self.right_hand.data.default_joint_pos)[env_ids] - delta_min = self.hand_dof_lower_limits[env_ids] - wp.to_torch(self.right_hand.data.default_joint_pos)[env_ids] + delta_max = self.hand_dof_upper_limits[env_ids] - self.right_hand.data.default_joint_pos.torch[env_ids] + delta_min = self.hand_dof_lower_limits[env_ids] - self.right_hand.data.default_joint_pos.torch[env_ids] dof_pos_noise = sample_uniform(-1.0, 1.0, (len(env_ids), self.num_hand_dofs), device=self.device) rand_delta = delta_min + (delta_max - delta_min) * 0.5 * dof_pos_noise - dof_pos = ( - wp.to_torch(self.right_hand.data.default_joint_pos)[env_ids] + self.cfg.reset_dof_pos_noise * rand_delta - ) + dof_pos = self.right_hand.data.default_joint_pos.torch[env_ids] + self.cfg.reset_dof_pos_noise * rand_delta dof_vel_noise = sample_uniform(-1.0, 1.0, (len(env_ids), self.num_hand_dofs), device=self.device) - dof_vel = ( - wp.to_torch(self.right_hand.data.default_joint_vel)[env_ids] + self.cfg.reset_dof_vel_noise * dof_vel_noise - ) + dof_vel = self.right_hand.data.default_joint_vel.torch[env_ids] + self.cfg.reset_dof_vel_noise * dof_vel_noise self.right_hand_prev_targets[env_ids] = dof_pos self.right_hand_curr_targets[env_ids] = dof_pos @@ -352,19 +348,15 @@ def _reset_idx(self, env_ids: Sequence[int] | torch.Tensor | None): self.right_hand.write_joint_velocity_to_sim_index(velocity=dof_vel, env_ids=env_ids) # reset left hand - delta_max = self.hand_dof_upper_limits[env_ids] - wp.to_torch(self.left_hand.data.default_joint_pos)[env_ids] - delta_min = self.hand_dof_lower_limits[env_ids] - wp.to_torch(self.left_hand.data.default_joint_pos)[env_ids] + delta_max = self.hand_dof_upper_limits[env_ids] - self.left_hand.data.default_joint_pos.torch[env_ids] + delta_min = self.hand_dof_lower_limits[env_ids] - self.left_hand.data.default_joint_pos.torch[env_ids] dof_pos_noise = sample_uniform(-1.0, 1.0, (len(env_ids), self.num_hand_dofs), device=self.device) rand_delta = delta_min + (delta_max - delta_min) * 0.5 * dof_pos_noise - dof_pos = ( - wp.to_torch(self.left_hand.data.default_joint_pos)[env_ids] + self.cfg.reset_dof_pos_noise * rand_delta - ) + dof_pos = self.left_hand.data.default_joint_pos.torch[env_ids] + self.cfg.reset_dof_pos_noise * rand_delta dof_vel_noise = sample_uniform(-1.0, 1.0, (len(env_ids), self.num_hand_dofs), device=self.device) - dof_vel = ( - wp.to_torch(self.left_hand.data.default_joint_vel)[env_ids] + self.cfg.reset_dof_vel_noise * dof_vel_noise - ) + dof_vel = self.left_hand.data.default_joint_vel.torch[env_ids] + self.cfg.reset_dof_vel_noise * dof_vel_noise self.left_hand_prev_targets[env_ids] = dof_pos self.left_hand_curr_targets[env_ids] = dof_pos @@ -390,33 +382,33 @@ def _reset_target_pose(self, env_ids): def _compute_intermediate_values(self): # data for right hand - self.right_fingertip_pos = wp.to_torch(self.right_hand.data.body_pos_w)[:, self.finger_bodies] - self.right_fingertip_rot = wp.to_torch(self.right_hand.data.body_quat_w)[:, self.finger_bodies] + self.right_fingertip_pos = self.right_hand.data.body_pos_w.torch[:, self.finger_bodies] + self.right_fingertip_rot = self.right_hand.data.body_quat_w.torch[:, self.finger_bodies] self.right_fingertip_pos -= self.scene.env_origins.repeat((1, self.num_fingertips)).reshape( self.num_envs, self.num_fingertips, 3 ) - self.right_fingertip_velocities = wp.to_torch(self.right_hand.data.body_vel_w)[:, self.finger_bodies] + self.right_fingertip_velocities = self.right_hand.data.body_vel_w.torch[:, self.finger_bodies] - self.right_hand_dof_pos = wp.to_torch(self.right_hand.data.joint_pos) - self.right_hand_dof_vel = wp.to_torch(self.right_hand.data.joint_vel) + self.right_hand_dof_pos = self.right_hand.data.joint_pos.torch + self.right_hand_dof_vel = self.right_hand.data.joint_vel.torch # data for left hand - self.left_fingertip_pos = wp.to_torch(self.left_hand.data.body_pos_w)[:, self.finger_bodies] - self.left_fingertip_rot = wp.to_torch(self.left_hand.data.body_quat_w)[:, self.finger_bodies] + self.left_fingertip_pos = self.left_hand.data.body_pos_w.torch[:, self.finger_bodies] + self.left_fingertip_rot = self.left_hand.data.body_quat_w.torch[:, self.finger_bodies] self.left_fingertip_pos -= self.scene.env_origins.repeat((1, self.num_fingertips)).reshape( self.num_envs, self.num_fingertips, 3 ) - self.left_fingertip_velocities = wp.to_torch(self.left_hand.data.body_vel_w)[:, self.finger_bodies] + self.left_fingertip_velocities = self.left_hand.data.body_vel_w.torch[:, self.finger_bodies] - self.left_hand_dof_pos = wp.to_torch(self.left_hand.data.joint_pos) - self.left_hand_dof_vel = wp.to_torch(self.left_hand.data.joint_vel) + self.left_hand_dof_pos = self.left_hand.data.joint_pos.torch + self.left_hand_dof_vel = self.left_hand.data.joint_vel.torch # data for object - self.object_pos = wp.to_torch(self.object.data.root_pos_w) - self.scene.env_origins - self.object_rot = wp.to_torch(self.object.data.root_quat_w) - self.object_velocities = wp.to_torch(self.object.data.root_vel_w) - self.object_linvel = wp.to_torch(self.object.data.root_lin_vel_w) - self.object_angvel = wp.to_torch(self.object.data.root_ang_vel_w) + self.object_pos = self.object.data.root_pos_w.torch - self.scene.env_origins + self.object_rot = self.object.data.root_quat_w.torch + self.object_velocities = self.object.data.root_vel_w.torch + self.object_linvel = self.object.data.root_lin_vel_w.torch + self.object_angvel = self.object.data.root_ang_vel_w.torch @torch.jit.script diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/classic/cartpole/mdp/rewards.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/classic/cartpole/mdp/rewards.py index 3bdb71d5606a..065a54e5b0b4 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/classic/cartpole/mdp/rewards.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/classic/cartpole/mdp/rewards.py @@ -8,7 +8,6 @@ from typing import TYPE_CHECKING import torch -import warp as wp from isaaclab.managers import SceneEntityCfg from isaaclab.utils.math import wrap_to_pi @@ -23,6 +22,6 @@ def joint_pos_target_l2(env: ManagerBasedRLEnv, target: float, asset_cfg: SceneE # extract the used quantities (to enable type-hinting) asset: Articulation = env.scene[asset_cfg.name] # wrap the joint positions to (-pi, pi) - joint_pos = wrap_to_pi(wp.to_torch(asset.data.joint_pos)[:, asset_cfg.joint_ids]) + joint_pos = wrap_to_pi(asset.data.joint_pos.torch[:, asset_cfg.joint_ids]) # compute the reward return torch.sum(torch.square(joint_pos - target), dim=1) diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/classic/humanoid/mdp/observations.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/classic/humanoid/mdp/observations.py index 5445328cc175..0c696ef62f50 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/classic/humanoid/mdp/observations.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/classic/humanoid/mdp/observations.py @@ -8,7 +8,6 @@ from typing import TYPE_CHECKING import torch -import warp as wp import isaaclab.utils.math as math_utils from isaaclab.managers import SceneEntityCfg @@ -23,7 +22,7 @@ def base_yaw_roll(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityC # extract the used quantities (to enable type-hinting) asset: Articulation = env.scene[asset_cfg.name] # extract euler angles (in world frame) - roll, _, yaw = math_utils.euler_xyz_from_quat(wp.to_torch(asset.data.root_quat_w)) + roll, _, yaw = math_utils.euler_xyz_from_quat(asset.data.root_quat_w.torch) # normalize angle to [-pi, pi] roll = torch.atan2(torch.sin(roll), torch.cos(roll)) yaw = torch.atan2(torch.sin(yaw), torch.cos(yaw)) @@ -36,7 +35,7 @@ def base_up_proj(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCf # extract the used quantities (to enable type-hinting) asset: Articulation = env.scene[asset_cfg.name] # compute base up vector - base_up_vec = -wp.to_torch(asset.data.projected_gravity_b) + base_up_vec = -asset.data.projected_gravity_b.torch return base_up_vec[:, 2].unsqueeze(-1) @@ -48,11 +47,11 @@ def base_heading_proj( # extract the used quantities (to enable type-hinting) asset: Articulation = env.scene[asset_cfg.name] # compute desired heading direction - to_target_pos = torch.tensor(target_pos, device=env.device) - wp.to_torch(asset.data.root_pos_w)[:, :3] + to_target_pos = torch.tensor(target_pos, device=env.device) - asset.data.root_pos_w.torch[:, :3] to_target_pos[:, 2] = 0.0 to_target_dir = math_utils.normalize(to_target_pos) # compute base forward vector - heading_vec = math_utils.quat_apply(wp.to_torch(asset.data.root_quat_w), wp.to_torch(asset.data.FORWARD_VEC_B)) + heading_vec = math_utils.quat_apply(asset.data.root_quat_w.torch, asset.data.FORWARD_VEC_B.torch) # compute dot product between heading and target direction heading_proj = torch.bmm(heading_vec.view(env.num_envs, 1, 3), to_target_dir.view(env.num_envs, 3, 1)) @@ -66,10 +65,10 @@ def base_angle_to_target( # extract the used quantities (to enable type-hinting) asset: Articulation = env.scene[asset_cfg.name] # compute desired heading direction - to_target_pos = torch.tensor(target_pos, device=env.device) - wp.to_torch(asset.data.root_pos_w)[:, :3] + to_target_pos = torch.tensor(target_pos, device=env.device) - asset.data.root_pos_w.torch[:, :3] walk_target_angle = torch.atan2(to_target_pos[:, 1], to_target_pos[:, 0]) # compute base forward vector - _, _, yaw = math_utils.euler_xyz_from_quat(wp.to_torch(asset.data.root_quat_w)) + _, _, yaw = math_utils.euler_xyz_from_quat(asset.data.root_quat_w.torch) # normalize angle to target to [-pi, pi] angle_to_target = walk_target_angle - yaw angle_to_target = torch.atan2(torch.sin(angle_to_target), torch.cos(angle_to_target)) diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/classic/humanoid/mdp/rewards.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/classic/humanoid/mdp/rewards.py index ceaa1371b917..dca907c50969 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/classic/humanoid/mdp/rewards.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/classic/humanoid/mdp/rewards.py @@ -8,7 +8,6 @@ from typing import TYPE_CHECKING import torch -import warp as wp import isaaclab.utils.math as math_utils import isaaclab.utils.string as string_utils @@ -55,7 +54,7 @@ def reset(self, env_ids: torch.Tensor): asset: Articulation = self._env.scene["robot"] # compute projection of current heading to desired heading vector target_pos = torch.tensor(self.cfg.params["target_pos"], device=self.device) - to_target_pos = target_pos - wp.to_torch(asset.data.root_pos_w)[env_ids, :3] + to_target_pos = target_pos - asset.data.root_pos_w.torch[env_ids, :3] # reward terms self.potentials[env_ids] = -torch.linalg.norm(to_target_pos, ord=2, dim=-1) / self._env.step_dt self.prev_potentials[env_ids] = self.potentials[env_ids] @@ -70,7 +69,7 @@ def __call__( asset: Articulation = env.scene[asset_cfg.name] # compute vector to target target_pos = torch.tensor(target_pos, device=env.device) - to_target_pos = target_pos - wp.to_torch(asset.data.root_pos_w)[:, :3] + to_target_pos = target_pos - asset.data.root_pos_w.torch[:, :3] to_target_pos[:, 2] = 0.0 # update history buffer and compute new potential self.prev_potentials[:] = self.potentials[:] @@ -107,9 +106,9 @@ def __call__( asset: Articulation = env.scene[asset_cfg.name] # compute the penalty over normalized joints joint_pos_scaled = math_utils.scale_transform( - wp.to_torch(asset.data.joint_pos), - wp.to_torch(asset.data.soft_joint_pos_limits)[..., 0], - wp.to_torch(asset.data.soft_joint_pos_limits)[..., 1], + asset.data.joint_pos.torch, + asset.data.soft_joint_pos_limits.torch[..., 0], + asset.data.soft_joint_pos_limits.torch[..., 1], ) # scale the violation amount by the gear ratio violation_amount = (torch.abs(joint_pos_scaled) - threshold) / (1 - threshold) @@ -145,5 +144,5 @@ def __call__( asset: Articulation = env.scene[asset_cfg.name] # return power = torque * velocity (here actions: joint torques) return torch.sum( - torch.abs(env.action_manager.action * wp.to_torch(asset.data.joint_vel) * self.gear_ratio_scaled), dim=-1 + torch.abs(env.action_manager.action * asset.data.joint_vel.torch * self.gear_ratio_scaled), dim=-1 ) diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/drone_arl/mdp/commands/drone_pose_command.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/drone_arl/mdp/commands/drone_pose_command.py index f3364c05c220..a9072367c83b 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/drone_arl/mdp/commands/drone_pose_command.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/drone_arl/mdp/commands/drone_pose_command.py @@ -8,7 +8,6 @@ from __future__ import annotations import torch -import warp as wp from isaaclab.envs.mdp.commands.pose_command import UniformPoseCommand from isaaclab.utils.math import combine_frame_transforms, compute_pose_error @@ -36,8 +35,8 @@ class DroneUniformPoseCommand(UniformPoseCommand): def _update_metrics(self): # transform command from base frame to simulation world frame self.pose_command_w[:, :3], self.pose_command_w[:, 3:] = combine_frame_transforms( - wp.to_torch(self.robot.data.root_pos_w), - wp.to_torch(self.robot.data.root_quat_w), + self.robot.data.root_pos_w.torch, + self.robot.data.root_quat_w.torch, self.pose_command_b[:, :3], self.pose_command_b[:, 3:], ) @@ -46,8 +45,8 @@ def _update_metrics(self): # Sub-terrain shift for correct position error calculation @grzemal self.pose_command_b[:, :3] + self._env.scene.env_origins, self.pose_command_w[:, 3:], - wp.to_torch(self.robot.data.body_pos_w)[:, self.body_idx], - wp.to_torch(self.robot.data.body_quat_w)[:, self.body_idx], + self.robot.data.body_pos_w.torch[:, self.body_idx], + self.robot.data.body_quat_w.torch[:, self.body_idx], ) self.metrics["position_error"] = torch.linalg.norm(pos_error, dim=-1) self.metrics["orientation_error"] = torch.linalg.norm(rot_error, dim=-1) @@ -64,5 +63,5 @@ def _debug_vis_callback(self, event): self.pose_command_b[:, :3] + self._env.scene.env_origins, self.pose_command_b[:, 3:] ) # -- current body pose - body_link_pose_w = wp.to_torch(self.robot.data.body_link_pose_w)[:, self.body_idx] + body_link_pose_w = self.robot.data.body_link_pose_w.torch[:, self.body_idx] self.current_pose_visualizer.visualize(body_link_pose_w[:, :3], body_link_pose_w[:, 3:7]) diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/drone_arl/mdp/observations.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/drone_arl/mdp/observations.py index 5b67ec4f66c3..c8b8048bd68b 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/drone_arl/mdp/observations.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/drone_arl/mdp/observations.py @@ -14,7 +14,6 @@ from typing import TYPE_CHECKING import torch -import warp as wp import isaaclab.utils.math as math_utils from isaaclab.managers import SceneEntityCfg @@ -50,7 +49,7 @@ def base_roll_pitch(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntit # extract the used quantities (to enable type-hinting) asset: Articulation = env.scene[asset_cfg.name] # extract euler angles (in world frame) - roll, pitch, _ = math_utils.euler_xyz_from_quat(wp.to_torch(asset.data.root_quat_w)) + roll, pitch, _ = math_utils.euler_xyz_from_quat(asset.data.root_quat_w.torch) # normalize angle to [-pi, pi] roll = math_utils.wrap_to_pi(roll) pitch = math_utils.wrap_to_pi(pitch) @@ -94,10 +93,10 @@ def generated_drone_commands( - A small epsilon (1e-8) is used to guard against zero-length direction vectors. """ asset: Multirotor = env.scene[asset_cfg.name] - current_position_w = wp.to_torch(asset.data.root_pos_w) - env.scene.env_origins + current_position_w = asset.data.root_pos_w.torch - env.scene.env_origins command = env.command_manager.get_command(command_name) current_position_b = math_utils.quat_apply_inverse( - wp.to_torch(asset.data.root_link_quat_w), command[:, :3] - current_position_w + asset.data.root_link_quat_w.torch, command[:, :3] - current_position_w ) current_position_b_dir = current_position_b / (torch.linalg.norm(current_position_b, dim=-1, keepdim=True) + 1e-8) current_position_b_mag = torch.linalg.norm(current_position_b, dim=-1, keepdim=True) diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/drone_arl/mdp/rewards.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/drone_arl/mdp/rewards.py index 840d694118d1..ce635cc544d8 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/drone_arl/mdp/rewards.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/drone_arl/mdp/rewards.py @@ -8,7 +8,6 @@ from typing import TYPE_CHECKING import torch -import warp as wp import isaaclab.utils.math as math_utils from isaaclab.managers import SceneEntityCfg @@ -51,7 +50,7 @@ def distance_to_goal_exp( command = env.command_manager.get_command(command_name) target_position_w = command[:, :3].clone() - current_position = wp.to_torch(asset.data.root_pos_w) - env.scene.env_origins + current_position = asset.data.root_pos_w.torch - env.scene.env_origins # compute the error position_error_square = torch.sum(torch.square(target_position_w - current_position), dim=1) @@ -82,7 +81,7 @@ def ang_vel_xyz_exp( asset: RigidObject = env.scene[asset_cfg.name] # compute squared magnitude of angular velocity (all axes) - ang_vel_squared = torch.sum(torch.square(wp.to_torch(asset.data.root_ang_vel_b)), dim=1) + ang_vel_squared = torch.sum(torch.square(asset.data.root_ang_vel_b.torch), dim=1) return torch.exp(-ang_vel_squared / std**2) @@ -109,7 +108,7 @@ def lin_vel_xyz_exp( asset: RigidObject = env.scene[asset_cfg.name] # compute squared magnitude of linear velocity (all axes) - lin_vel_squared = torch.sum(torch.square(wp.to_torch(asset.data.root_lin_vel_w)), dim=1) + lin_vel_squared = torch.sum(torch.square(asset.data.root_lin_vel_w.torch), dim=1) return torch.exp(-lin_vel_squared / std**2) @@ -138,7 +137,7 @@ def yaw_aligned( asset: RigidObject = env.scene[asset_cfg.name] # extract yaw from current orientation - _, _, yaw = math_utils.euler_xyz_from_quat(wp.to_torch(asset.data.root_quat_w)) + _, _, yaw = math_utils.euler_xyz_from_quat(asset.data.root_quat_w.torch) # normalize yaw to [-pi, pi] (target is 0) yaw = math_utils.wrap_to_pi(yaw) diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomanipulation/pick_place/mdp/actions.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomanipulation/pick_place/mdp/actions.py index 307dd026fdac..32c768981b3e 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomanipulation/pick_place/mdp/actions.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomanipulation/pick_place/mdp/actions.py @@ -8,7 +8,6 @@ from typing import TYPE_CHECKING import torch -import warp as wp from isaaclab.assets.articulation import Articulation from isaaclab.managers.action_manager import ActionTerm @@ -47,7 +46,7 @@ def __init__(self, cfg: AgileBasedLowerBodyActionCfg, env: ManagerBasedEnv): # Get the scale and offset from the configuration self._policy_output_scale = torch.tensor(cfg.policy_output_scale, device=env.device) - self._policy_output_offset = wp.to_torch(self._asset.data.default_joint_pos)[:, self._joint_ids].clone() + self._policy_output_offset = self._asset.data.default_joint_pos.torch[:, self._joint_ids].clone() # Create tensors to store raw and processed actions self._raw_actions = torch.zeros(self.num_envs, len(self._joint_ids), device=self.device) diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomanipulation/pick_place/mdp/observations.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomanipulation/pick_place/mdp/observations.py index 5e39567efa50..4057bc6fe590 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomanipulation/pick_place/mdp/observations.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomanipulation/pick_place/mdp/observations.py @@ -8,7 +8,6 @@ from typing import TYPE_CHECKING import torch -import warp as wp if TYPE_CHECKING: from isaaclab.envs import ManagerBasedRLEnv @@ -22,7 +21,7 @@ def upper_body_last_action( ) -> torch.Tensor: """Extract the last action of the upper body.""" asset = env.scene[asset_cfg.name] - joint_pos_target = wp.to_torch(asset.data.joint_pos_target) + joint_pos_target = asset.data.joint_pos_target.torch # Use joint_names from asset_cfg to find indices joint_names = asset_cfg.joint_names if hasattr(asset_cfg, "joint_names") else None diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomanipulation/pick_place/mdp/terminations.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomanipulation/pick_place/mdp/terminations.py index 8e530a7d71e0..db89776c3e37 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomanipulation/pick_place/mdp/terminations.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomanipulation/pick_place/mdp/terminations.py @@ -10,7 +10,6 @@ from typing import TYPE_CHECKING import torch -import warp as wp from isaaclab.managers import SceneEntityCfg from isaaclab.utils.math import quat_apply_inverse @@ -75,9 +74,10 @@ def task_done_pick_place_table_frame( # Get table world pose table_pos_w, table_quat_w = table.get_world_poses() + table_pos_w, table_quat_w = table_pos_w.torch, table_quat_w.torch # Broadcast table pose if a single table is shared across all envs - object_root_pos_w = wp.to_torch(object.data.root_pos_w) # [num_envs, 3] + object_root_pos_w = object.data.root_pos_w.torch # [num_envs, 3] if table_pos_w.shape[0] == 1 and object_root_pos_w.shape[0] > 1: table_pos_w = table_pos_w.expand(object_root_pos_w.shape[0], -1) table_quat_w = table_quat_w.expand(object_root_pos_w.shape[0], -1) @@ -87,10 +87,10 @@ def task_done_pick_place_table_frame( object_x = object_to_table_pos[:, 0] object_y = object_to_table_pos[:, 1] object_height = object_to_table_pos[:, 2] - object_vel = torch.abs(wp.to_torch(object.data.root_vel_w)) + object_vel = torch.abs(object.data.root_vel_w.torch) # Right wrist position in table frame - robot_body_pos_w = wp.to_torch(env.scene["robot"].data.body_pos_w) + robot_body_pos_w = env.scene["robot"].data.body_pos_w.torch right_eef_idx = env.scene["robot"].data.body_names.index(task_link_name) right_wrist_pos_w = robot_body_pos_w[:, right_eef_idx, :] - table_pos_w right_wrist_x = quat_apply_inverse(table_quat_w, right_wrist_pos_w)[:, 0] @@ -146,8 +146,8 @@ def object_too_far_from_robot( robot: RigidObject = env.scene[robot_cfg.name] object: RigidObject = env.scene[object_cfg.name] - robot_pos = wp.to_torch(robot.data.root_pos_w)[:, :3] - object_pos = wp.to_torch(object.data.root_pos_w)[:, :3] + robot_pos = robot.data.root_pos_w.torch[:, :3] + object_pos = object.data.root_pos_w.torch[:, :3] distance = torch.norm(robot_pos - object_pos, dim=1) diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/config/spot/mdp/events.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/config/spot/mdp/events.py index bd89176fc5d1..eb9929ada4e1 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/config/spot/mdp/events.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/config/spot/mdp/events.py @@ -15,7 +15,6 @@ from typing import TYPE_CHECKING import torch -import warp as wp from isaaclab.managers import SceneEntityCfg from isaaclab.utils.math import sample_uniform @@ -41,16 +40,16 @@ def reset_joints_around_default( # extract the used quantities (to enable type-hinting) asset: Articulation = env.scene[asset_cfg.name] # get default joint state - joint_min_pos = wp.to_torch(asset.data.default_joint_pos)[env_ids] + position_range[0] - joint_max_pos = wp.to_torch(asset.data.default_joint_pos)[env_ids] + position_range[1] - joint_min_vel = wp.to_torch(asset.data.default_joint_vel)[env_ids] + velocity_range[0] - joint_max_vel = wp.to_torch(asset.data.default_joint_vel)[env_ids] + velocity_range[1] + joint_min_pos = asset.data.default_joint_pos.torch[env_ids] + position_range[0] + joint_max_pos = asset.data.default_joint_pos.torch[env_ids] + position_range[1] + joint_min_vel = asset.data.default_joint_vel.torch[env_ids] + velocity_range[0] + joint_max_vel = asset.data.default_joint_vel.torch[env_ids] + velocity_range[1] # clip pos to range - joint_pos_limits = wp.to_torch(asset.data.soft_joint_pos_limits)[env_ids, ...] + joint_pos_limits = asset.data.soft_joint_pos_limits.torch[env_ids, ...] joint_min_pos = torch.clamp(joint_min_pos, min=joint_pos_limits[..., 0], max=joint_pos_limits[..., 1]) joint_max_pos = torch.clamp(joint_max_pos, min=joint_pos_limits[..., 0], max=joint_pos_limits[..., 1]) # clip vel to range - joint_vel_abs_limits = wp.to_torch(asset.data.soft_joint_vel_limits)[env_ids] + joint_vel_abs_limits = asset.data.soft_joint_vel_limits.torch[env_ids] joint_min_vel = torch.clamp(joint_min_vel, min=-joint_vel_abs_limits, max=joint_vel_abs_limits) joint_max_vel = torch.clamp(joint_max_vel, min=-joint_vel_abs_limits, max=joint_vel_abs_limits) # sample these values randomly diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/config/spot/mdp/rewards.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/config/spot/mdp/rewards.py index 81e72ee50727..5cb88e0420aa 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/config/spot/mdp/rewards.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/config/spot/mdp/rewards.py @@ -14,7 +14,6 @@ from typing import TYPE_CHECKING import torch -import warp as wp from isaaclab.managers import ManagerTermBase, SceneEntityCfg @@ -44,14 +43,14 @@ def air_time_reward( if contact_sensor.cfg.track_air_time is False: raise RuntimeError("Activate ContactSensor's track_air_time!") # compute the reward - current_air_time = wp.to_torch(contact_sensor.data.current_air_time)[:, sensor_cfg.body_ids] - current_contact_time = wp.to_torch(contact_sensor.data.current_contact_time)[:, sensor_cfg.body_ids] + current_air_time = contact_sensor.data.current_air_time.torch[:, sensor_cfg.body_ids] + current_contact_time = contact_sensor.data.current_contact_time.torch[:, sensor_cfg.body_ids] t_max = torch.max(current_air_time, current_contact_time) t_min = torch.clip(t_max, max=mode_time) stance_cmd_reward = torch.clip(current_contact_time - current_air_time, -mode_time, mode_time) cmd = torch.linalg.norm(env.command_manager.get_command("base_velocity"), dim=1).unsqueeze(dim=1).expand(-1, 4) - body_vel = torch.linalg.norm(wp.to_torch(asset.data.root_lin_vel_b)[:, :2], dim=1).unsqueeze(dim=1).expand(-1, 4) + body_vel = torch.linalg.norm(asset.data.root_lin_vel_b.torch[:, :2], dim=1).unsqueeze(dim=1).expand(-1, 4) reward = torch.where( torch.logical_or(cmd > 0.0, body_vel > velocity_threshold), torch.where(t_max < mode_time, t_min, 0), @@ -66,7 +65,7 @@ def base_angular_velocity_reward(env: ManagerBasedRLEnv, asset_cfg: SceneEntityC asset: RigidObject = env.scene[asset_cfg.name] # compute the error target = env.command_manager.get_command("base_velocity")[:, 2] - ang_vel_error = torch.linalg.norm((target - wp.to_torch(asset.data.root_ang_vel_b)[:, 2]).unsqueeze(1), dim=1) + ang_vel_error = torch.linalg.norm((target - asset.data.root_ang_vel_b.torch[:, 2]).unsqueeze(1), dim=1) return torch.exp(-ang_vel_error / std) @@ -78,7 +77,7 @@ def base_linear_velocity_reward( asset: RigidObject = env.scene[asset_cfg.name] # compute the error target = env.command_manager.get_command("base_velocity")[:, :2] - lin_vel_error = torch.linalg.norm((target - wp.to_torch(asset.data.root_lin_vel_b)[:, :2]), dim=1) + lin_vel_error = torch.linalg.norm((target - asset.data.root_lin_vel_b.torch[:, :2]), dim=1) # fixed 1.0 multiple for tracking below the ramp_at_vel value, then scale by the rate above vel_cmd_magnitude = torch.linalg.norm(target, dim=1) velocity_scaling_multiple = torch.clamp(1.0 + ramp_rate * (vel_cmd_magnitude - ramp_at_vel), min=1.0) @@ -151,7 +150,7 @@ def __call__( async_reward = async_reward_0 * async_reward_1 * async_reward_2 * async_reward_3 # only enforce gait if cmd > 0 cmd = torch.linalg.norm(env.command_manager.get_command("base_velocity"), dim=1) - body_vel = torch.linalg.norm(wp.to_torch(self.asset.data.root_lin_vel_b)[:, :2], dim=1) + body_vel = torch.linalg.norm(self.asset.data.root_lin_vel_b.torch[:, :2], dim=1) return torch.where( torch.logical_or(cmd > 0.0, body_vel > self.velocity_threshold), sync_reward * async_reward, 0.0 ) @@ -162,8 +161,8 @@ def __call__( def _sync_reward_func(self, foot_0: int, foot_1: int) -> torch.Tensor: """Reward synchronization of two feet.""" - air_time = wp.to_torch(self.contact_sensor.data.current_air_time) - contact_time = wp.to_torch(self.contact_sensor.data.current_contact_time) + air_time = self.contact_sensor.data.current_air_time.torch + contact_time = self.contact_sensor.data.current_contact_time.torch # penalize the difference between the most recent air time and contact time of synced feet pairs. se_air = torch.clip(torch.square(air_time[:, foot_0] - air_time[:, foot_1]), max=self.max_err**2) se_contact = torch.clip(torch.square(contact_time[:, foot_0] - contact_time[:, foot_1]), max=self.max_err**2) @@ -171,8 +170,8 @@ def _sync_reward_func(self, foot_0: int, foot_1: int) -> torch.Tensor: def _async_reward_func(self, foot_0: int, foot_1: int) -> torch.Tensor: """Reward anti-synchronization of two feet.""" - air_time = wp.to_torch(self.contact_sensor.data.current_air_time) - contact_time = wp.to_torch(self.contact_sensor.data.current_contact_time) + air_time = self.contact_sensor.data.current_air_time.torch + contact_time = self.contact_sensor.data.current_contact_time.torch # penalize the difference between opposing contact modes air time of feet 1 to contact time of feet 2 # and contact time of feet 1 to air time of feet 2) of feet pairs that are not in sync with each other. se_act_0 = torch.clip(torch.square(air_time[:, foot_0] - contact_time[:, foot_1]), max=self.max_err**2) @@ -185,9 +184,9 @@ def foot_clearance_reward( ) -> torch.Tensor: """Reward the swinging feet for clearing a specified height off the ground""" asset: RigidObject = env.scene[asset_cfg.name] - foot_z_target_error = torch.square(wp.to_torch(asset.data.body_pos_w)[:, asset_cfg.body_ids, 2] - target_height) + foot_z_target_error = torch.square(asset.data.body_pos_w.torch[:, asset_cfg.body_ids, 2] - target_height) foot_velocity_tanh = torch.tanh( - tanh_mult * torch.linalg.norm(wp.to_torch(asset.data.body_lin_vel_w)[:, asset_cfg.body_ids, :2], dim=2) + tanh_mult * torch.linalg.norm(asset.data.body_lin_vel_w.torch[:, asset_cfg.body_ids, :2], dim=2) ) reward = foot_z_target_error * foot_velocity_tanh return torch.exp(-torch.sum(reward, dim=1) / std) @@ -210,8 +209,8 @@ def air_time_variance_penalty(env: ManagerBasedRLEnv, sensor_cfg: SceneEntityCfg if contact_sensor.cfg.track_air_time is False: raise RuntimeError("Activate ContactSensor's track_air_time!") # compute the reward - last_air_time = wp.to_torch(contact_sensor.data.last_air_time)[:, sensor_cfg.body_ids] - last_contact_time = wp.to_torch(contact_sensor.data.last_contact_time)[:, sensor_cfg.body_ids] + last_air_time = contact_sensor.data.last_air_time.torch[:, sensor_cfg.body_ids] + last_contact_time = contact_sensor.data.last_contact_time.torch[:, sensor_cfg.body_ids] return torch.var(torch.clip(last_air_time, max=0.5), dim=1) + torch.var( torch.clip(last_contact_time, max=0.5), dim=1 ) @@ -222,8 +221,8 @@ def base_motion_penalty(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg) -> to """Penalize base vertical and roll/pitch velocity""" # extract the used quantities (to enable type-hinting) asset: RigidObject = env.scene[asset_cfg.name] - return 0.8 * torch.square(wp.to_torch(asset.data.root_lin_vel_b)[:, 2]) + 0.2 * torch.sum( - torch.abs(wp.to_torch(asset.data.root_ang_vel_b)[:, :2]), dim=1 + return 0.8 * torch.square(asset.data.root_lin_vel_b.torch[:, 2]) + 0.2 * torch.sum( + torch.abs(asset.data.root_ang_vel_b.torch[:, :2]), dim=1 ) @@ -234,7 +233,7 @@ def base_orientation_penalty(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg) """ # extract the used quantities (to enable type-hinting) asset: RigidObject = env.scene[asset_cfg.name] - return torch.linalg.norm((wp.to_torch(asset.data.projected_gravity_b)[:, :2]), dim=1) + return torch.linalg.norm((asset.data.projected_gravity_b.torch[:, :2]), dim=1) def foot_slip_penalty( @@ -246,11 +245,11 @@ def foot_slip_penalty( contact_sensor: ContactSensor = env.scene.sensors[sensor_cfg.name] # check if contact force is above threshold - net_contact_forces = wp.to_torch(contact_sensor.data.net_forces_w_history) + net_contact_forces = contact_sensor.data.net_forces_w_history.torch is_contact = ( torch.max(torch.linalg.norm(net_contact_forces[:, :, sensor_cfg.body_ids], dim=-1), dim=1)[0] > threshold ) - foot_planar_velocity = torch.linalg.norm(wp.to_torch(asset.data.body_lin_vel_w)[:, asset_cfg.body_ids, :2], dim=2) + foot_planar_velocity = torch.linalg.norm(asset.data.body_lin_vel_w.torch[:, asset_cfg.body_ids, :2], dim=2) reward = is_contact * foot_planar_velocity return torch.sum(reward, dim=1) @@ -260,7 +259,7 @@ def joint_acceleration_penalty(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg """Penalize joint accelerations on the articulation.""" # extract the used quantities (to enable type-hinting) asset: Articulation = env.scene[asset_cfg.name] - return torch.linalg.norm((wp.to_torch(asset.data.joint_acc)), dim=1) + return torch.linalg.norm((asset.data.joint_acc.torch), dim=1) def joint_position_penalty( @@ -270,8 +269,8 @@ def joint_position_penalty( # extract the used quantities (to enable type-hinting) asset: Articulation = env.scene[asset_cfg.name] cmd = torch.linalg.norm(env.command_manager.get_command("base_velocity"), dim=1) - body_vel = torch.linalg.norm(wp.to_torch(asset.data.root_lin_vel_b)[:, :2], dim=1) - reward = torch.linalg.norm((wp.to_torch(asset.data.joint_pos) - wp.to_torch(asset.data.default_joint_pos)), dim=1) + body_vel = torch.linalg.norm(asset.data.root_lin_vel_b.torch[:, :2], dim=1) + reward = torch.linalg.norm((asset.data.joint_pos.torch - asset.data.default_joint_pos.torch), dim=1) return torch.where(torch.logical_or(cmd > 0.0, body_vel > velocity_threshold), reward, stand_still_scale * reward) @@ -279,11 +278,11 @@ def joint_torques_penalty(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg) -> """Penalize joint torques on the articulation.""" # extract the used quantities (to enable type-hinting) asset: Articulation = env.scene[asset_cfg.name] - return torch.linalg.norm((wp.to_torch(asset.data.applied_torque)), dim=1) + return torch.linalg.norm((asset.data.applied_torque.torch), dim=1) def joint_velocity_penalty(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg) -> torch.Tensor: """Penalize joint velocities on the articulation.""" # extract the used quantities (to enable type-hinting) asset: Articulation = env.scene[asset_cfg.name] - return torch.linalg.norm((wp.to_torch(asset.data.joint_vel)), dim=1) + return torch.linalg.norm((asset.data.joint_vel.torch), dim=1) diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/mdp/curriculums.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/mdp/curriculums.py index ece0a6161150..1438cbf03801 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/mdp/curriculums.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/mdp/curriculums.py @@ -15,7 +15,6 @@ from typing import TYPE_CHECKING import torch -import warp as wp from isaaclab.managers import SceneEntityCfg @@ -45,9 +44,7 @@ def terrain_levels_vel( terrain: TerrainImporter = env.scene.terrain command = env.command_manager.get_command("base_velocity") # compute the distance the robot walked - distance = torch.linalg.norm( - wp.to_torch(asset.data.root_pos_w)[env_ids, :2] - env.scene.env_origins[env_ids, :2], dim=1 - ) + distance = torch.linalg.norm(asset.data.root_pos_w.torch[env_ids, :2] - env.scene.env_origins[env_ids, :2], dim=1) # robots that walked far enough progress to harder terrains move_up = distance > terrain.cfg.terrain_generator.size[0] / 2 # robots that walked less than half of their required distance go to simpler terrains diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/mdp/rewards.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/mdp/rewards.py index 2198cff0cead..a5168f6c95d9 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/mdp/rewards.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/mdp/rewards.py @@ -14,7 +14,6 @@ from typing import TYPE_CHECKING import torch -import warp as wp from isaaclab.envs import mdp from isaaclab.managers import SceneEntityCfg @@ -39,8 +38,8 @@ def feet_air_time( # extract the used quantities (to enable type-hinting) contact_sensor: ContactSensor = env.scene.sensors[sensor_cfg.name] # compute the reward - first_contact = wp.to_torch(contact_sensor.compute_first_contact(env.step_dt))[:, sensor_cfg.body_ids] - last_air_time = wp.to_torch(contact_sensor.data.last_air_time)[:, sensor_cfg.body_ids] + first_contact = contact_sensor.compute_first_contact(env.step_dt).torch[:, sensor_cfg.body_ids] + last_air_time = contact_sensor.data.last_air_time.torch[:, sensor_cfg.body_ids] reward = torch.sum((last_air_time - threshold) * first_contact, dim=1) # no reward for zero command reward *= torch.linalg.norm(env.command_manager.get_command(command_name)[:, :2], dim=1) > 0.1 @@ -57,8 +56,8 @@ def feet_air_time_positive_biped(env, command_name: str, threshold: float, senso """ contact_sensor: ContactSensor = env.scene.sensors[sensor_cfg.name] # compute the reward - air_time = wp.to_torch(contact_sensor.data.current_air_time)[:, sensor_cfg.body_ids] - contact_time = wp.to_torch(contact_sensor.data.current_contact_time)[:, sensor_cfg.body_ids] + air_time = contact_sensor.data.current_air_time.torch[:, sensor_cfg.body_ids] + contact_time = contact_sensor.data.current_contact_time.torch[:, sensor_cfg.body_ids] in_contact = contact_time > 0.0 in_mode_time = torch.where(in_contact, contact_time, air_time) single_stance = torch.sum(in_contact.int(), dim=1) == 1 @@ -79,12 +78,11 @@ def feet_slide(env, sensor_cfg: SceneEntityCfg, asset_cfg: SceneEntityCfg = Scen # Penalize feet sliding contact_sensor: ContactSensor = env.scene.sensors[sensor_cfg.name] contacts = ( - wp.to_torch(contact_sensor.data.net_forces_w_history)[:, :, sensor_cfg.body_ids, :].norm(dim=-1).max(dim=1)[0] - > 1.0 + contact_sensor.data.net_forces_w_history.torch[:, :, sensor_cfg.body_ids, :].norm(dim=-1).max(dim=1)[0] > 1.0 ) asset = env.scene[asset_cfg.name] - body_vel = wp.to_torch(asset.data.body_lin_vel_w)[:, asset_cfg.body_ids, :2] + body_vel = asset.data.body_lin_vel_w.torch[:, asset_cfg.body_ids, :2] reward = torch.sum(body_vel.norm(dim=-1) * contacts, dim=1) return reward @@ -97,9 +95,7 @@ def track_lin_vel_xy_yaw_frame_exp( """ # extract the used quantities (to enable type-hinting) asset = env.scene[asset_cfg.name] - vel_yaw = quat_apply_inverse( - yaw_quat(wp.to_torch(asset.data.root_quat_w)), wp.to_torch(asset.data.root_lin_vel_w)[:, :3] - ) + vel_yaw = quat_apply_inverse(yaw_quat(asset.data.root_quat_w.torch), asset.data.root_lin_vel_w.torch[:, :3]) lin_vel_error = torch.sum( torch.square(env.command_manager.get_command(command_name)[:, :2] - vel_yaw[:, :2]), dim=1 ) @@ -113,7 +109,7 @@ def track_ang_vel_z_world_exp( # extract the used quantities (to enable type-hinting) asset = env.scene[asset_cfg.name] ang_vel_error = torch.square( - env.command_manager.get_command(command_name)[:, 2] - wp.to_torch(asset.data.root_ang_vel_w)[:, 2] + env.command_manager.get_command(command_name)[:, 2] - asset.data.root_ang_vel_w.torch[:, 2] ) return torch.exp(-ang_vel_error / std**2) diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/mdp/terminations.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/mdp/terminations.py index 63add623c465..5559250d9555 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/mdp/terminations.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/mdp/terminations.py @@ -14,7 +14,6 @@ from typing import TYPE_CHECKING import torch -import warp as wp from isaaclab.managers import SceneEntityCfg @@ -48,8 +47,8 @@ def terrain_out_of_bounds( asset: RigidObject = env.scene[asset_cfg.name] # check if the agent is out of bounds - x_out_of_bounds = torch.abs(wp.to_torch(asset.data.root_pos_w)[:, 0]) > 0.5 * map_width - distance_buffer - y_out_of_bounds = torch.abs(wp.to_torch(asset.data.root_pos_w)[:, 1]) > 0.5 * map_height - distance_buffer + x_out_of_bounds = torch.abs(asset.data.root_pos_w.torch[:, 0]) > 0.5 * map_width - distance_buffer + y_out_of_bounds = torch.abs(asset.data.root_pos_w.torch[:, 1]) > 0.5 * map_height - distance_buffer return torch.logical_or(x_out_of_bounds, y_out_of_bounds) else: raise ValueError("Received unsupported terrain type, must be either 'plane' or 'generator'.") diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/cabinet/mdp/observations.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/cabinet/mdp/observations.py index 75e1ff44157f..1da8557cc0ca 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/cabinet/mdp/observations.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/cabinet/mdp/observations.py @@ -8,7 +8,6 @@ from typing import TYPE_CHECKING import torch -import warp as wp import isaaclab.utils.math as math_utils @@ -23,7 +22,7 @@ def rel_ee_object_distance(env: ManagerBasedRLEnv) -> torch.Tensor: ee_tf_data: FrameTransformerData = env.scene["ee_frame"].data object_data: ArticulationData = env.scene["object"].data - return wp.to_torch(object_data.root_pos_w) - wp.to_torch(ee_tf_data.target_pos_w)[..., 0, :] + return object_data.root_pos_w.torch - ee_tf_data.target_pos_w.torch[..., 0, :] def rel_ee_drawer_distance(env: ManagerBasedRLEnv) -> torch.Tensor: @@ -31,13 +30,13 @@ def rel_ee_drawer_distance(env: ManagerBasedRLEnv) -> torch.Tensor: ee_tf_data: FrameTransformerData = env.scene["ee_frame"].data cabinet_tf_data: FrameTransformerData = env.scene["cabinet_frame"].data - return wp.to_torch(cabinet_tf_data.target_pos_w)[..., 0, :] - wp.to_torch(ee_tf_data.target_pos_w)[..., 0, :] + return cabinet_tf_data.target_pos_w.torch[..., 0, :] - ee_tf_data.target_pos_w.torch[..., 0, :] def fingertips_pos(env: ManagerBasedRLEnv) -> torch.Tensor: """The position of the fingertips relative to the environment origins.""" ee_tf_data: FrameTransformerData = env.scene["ee_frame"].data - fingertips_pos = wp.to_torch(ee_tf_data.target_pos_w)[..., 1:, :] - env.scene.env_origins.unsqueeze(1) + fingertips_pos = ee_tf_data.target_pos_w.torch[..., 1:, :] - env.scene.env_origins.unsqueeze(1) return fingertips_pos.view(env.num_envs, -1) @@ -45,7 +44,7 @@ def fingertips_pos(env: ManagerBasedRLEnv) -> torch.Tensor: def ee_pos(env: ManagerBasedRLEnv) -> torch.Tensor: """The position of the end-effector relative to the environment origins.""" ee_tf_data: FrameTransformerData = env.scene["ee_frame"].data - ee_pos = wp.to_torch(ee_tf_data.target_pos_w)[..., 0, :] - env.scene.env_origins + ee_pos = ee_tf_data.target_pos_w.torch[..., 0, :] - env.scene.env_origins return ee_pos @@ -56,6 +55,6 @@ def ee_quat(env: ManagerBasedRLEnv, make_quat_unique: bool = True) -> torch.Tens If :attr:`make_quat_unique` is True, the quaternion is made unique by ensuring the real part is positive. """ ee_tf_data: FrameTransformerData = env.scene["ee_frame"].data - ee_quat = wp.to_torch(ee_tf_data.target_quat_w)[..., 0, :] + ee_quat = ee_tf_data.target_quat_w.torch[..., 0, :] # make first element of quaternion positive return math_utils.quat_unique(ee_quat) if make_quat_unique else ee_quat diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/cabinet/mdp/rewards.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/cabinet/mdp/rewards.py index aeb87e29e224..dfbfeb4b0846 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/cabinet/mdp/rewards.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/cabinet/mdp/rewards.py @@ -8,7 +8,6 @@ from typing import TYPE_CHECKING import torch -import warp as wp from isaaclab.managers import SceneEntityCfg from isaaclab.utils.math import matrix_from_quat @@ -30,8 +29,8 @@ def approach_ee_handle(env: ManagerBasedRLEnv, threshold: float) -> torch.Tensor \end{cases} """ - ee_tcp_pos = wp.to_torch(env.scene["ee_frame"].data.target_pos_w)[..., 0, :] - handle_pos = wp.to_torch(env.scene["cabinet_frame"].data.target_pos_w)[..., 0, :] + ee_tcp_pos = env.scene["ee_frame"].data.target_pos_w.torch[..., 0, :] + handle_pos = env.scene["cabinet_frame"].data.target_pos_w.torch[..., 0, :] # Compute the distance of the end-effector to the handle distance = torch.linalg.norm(handle_pos - ee_tcp_pos, dim=-1, ord=2) @@ -54,8 +53,8 @@ def align_ee_handle(env: ManagerBasedRLEnv) -> torch.Tensor: where :math:`align_z` is the dot product of the z direction of the gripper and the -x direction of the handle and :math:`align_x` is the dot product of the x direction of the gripper and the -y direction of the handle. """ - ee_tcp_quat = wp.to_torch(env.scene["ee_frame"].data.target_quat_w)[..., 0, :] - handle_quat = wp.to_torch(env.scene["cabinet_frame"].data.target_quat_w)[..., 0, :] + ee_tcp_quat = env.scene["ee_frame"].data.target_quat_w.torch[..., 0, :] + handle_quat = env.scene["cabinet_frame"].data.target_quat_w.torch[..., 0, :] ee_tcp_rot_mat = matrix_from_quat(ee_tcp_quat) handle_mat = matrix_from_quat(handle_quat) @@ -80,9 +79,9 @@ def align_grasp_around_handle(env: ManagerBasedRLEnv) -> torch.Tensor: The correct hand orientation is when the left finger is above the handle and the right finger is below the handle. """ # Target object position: (num_envs, 3) - handle_pos = wp.to_torch(env.scene["cabinet_frame"].data.target_pos_w)[..., 0, :] + handle_pos = env.scene["cabinet_frame"].data.target_pos_w.torch[..., 0, :] # Fingertips position: (num_envs, n_fingertips, 3) - ee_fingertips_w = wp.to_torch(env.scene["ee_frame"].data.target_pos_w)[..., 1:, :] + ee_fingertips_w = env.scene["ee_frame"].data.target_pos_w.torch[..., 1:, :] lfinger_pos = ee_fingertips_w[..., 0, :] rfinger_pos = ee_fingertips_w[..., 1, :] @@ -100,9 +99,9 @@ def approach_gripper_handle(env: ManagerBasedRLEnv, offset: float = 0.04) -> tor (i.e., the left finger is above the handle and the right finger is below the handle). Otherwise, it returns zero. """ # Target object position: (num_envs, 3) - handle_pos = wp.to_torch(env.scene["cabinet_frame"].data.target_pos_w)[..., 0, :] + handle_pos = env.scene["cabinet_frame"].data.target_pos_w.torch[..., 0, :] # Fingertips position: (num_envs, n_fingertips, 3) - ee_fingertips_w = wp.to_torch(env.scene["ee_frame"].data.target_pos_w)[..., 1:, :] + ee_fingertips_w = env.scene["ee_frame"].data.target_pos_w.torch[..., 1:, :] lfinger_pos = ee_fingertips_w[..., 0, :] rfinger_pos = ee_fingertips_w[..., 1, :] @@ -127,9 +126,9 @@ def grasp_handle( Note: It is assumed that zero joint position corresponds to the fingers being closed. """ - ee_tcp_pos = wp.to_torch(env.scene["ee_frame"].data.target_pos_w)[..., 0, :] - handle_pos = wp.to_torch(env.scene["cabinet_frame"].data.target_pos_w)[..., 0, :] - gripper_joint_pos = wp.to_torch(env.scene[asset_cfg.name].data.joint_pos)[:, asset_cfg.joint_ids] + ee_tcp_pos = env.scene["ee_frame"].data.target_pos_w.torch[..., 0, :] + handle_pos = env.scene["cabinet_frame"].data.target_pos_w.torch[..., 0, :] + gripper_joint_pos = env.scene[asset_cfg.name].data.joint_pos.torch[:, asset_cfg.joint_ids] distance = torch.linalg.norm(handle_pos - ee_tcp_pos, dim=-1, ord=2) is_close = distance <= threshold @@ -142,7 +141,7 @@ def open_drawer_bonus(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg) -> torc The bonus is given when the drawer is open. If the grasp is around the handle, the bonus is doubled. """ - drawer_pos = wp.to_torch(env.scene[asset_cfg.name].data.joint_pos)[:, asset_cfg.joint_ids[0]] + drawer_pos = env.scene[asset_cfg.name].data.joint_pos.torch[:, asset_cfg.joint_ids[0]] is_graspable = align_grasp_around_handle(env).float() return (is_graspable + 1.0) * drawer_pos @@ -154,7 +153,7 @@ def multi_stage_open_drawer(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg) - Depending on the drawer's position, the reward is given in three stages: easy, medium, and hard. This helps the agent to learn to open the drawer in a controlled manner. """ - drawer_pos = wp.to_torch(env.scene[asset_cfg.name].data.joint_pos)[:, asset_cfg.joint_ids[0]] + drawer_pos = env.scene[asset_cfg.name].data.joint_pos.torch[:, asset_cfg.joint_ids[0]] is_graspable = align_grasp_around_handle(env).float() open_easy = (drawer_pos > 0.01) * 0.5 diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/deploy/mdp/events.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/deploy/mdp/events.py index 503fd3af75d8..2933033fd8e0 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/deploy/mdp/events.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/deploy/mdp/events.py @@ -254,24 +254,24 @@ def __call__( # IK loop for i in range(max_iterations): # Get current joint state - joint_pos = wp.to_torch(self.robot_asset.data.joint_pos)[env_ids].clone() - joint_vel = wp.to_torch(self.robot_asset.data.joint_vel)[env_ids].clone() + joint_pos = self.robot_asset.data.joint_pos.torch[env_ids].clone() + joint_vel = self.robot_asset.data.joint_vel.torch[env_ids].clone() # Stack all gear positions and quaternions all_gear_pos = torch.stack( [ - wp.to_torch(env.scene["factory_gear_small"].data.root_link_pos_w), - wp.to_torch(env.scene["factory_gear_medium"].data.root_link_pos_w), - wp.to_torch(env.scene["factory_gear_large"].data.root_link_pos_w), + env.scene["factory_gear_small"].data.root_link_pos_w.torch, + env.scene["factory_gear_medium"].data.root_link_pos_w.torch, + env.scene["factory_gear_large"].data.root_link_pos_w.torch, ], dim=1, )[env_ids] all_gear_quat = torch.stack( [ - wp.to_torch(env.scene["factory_gear_small"].data.root_link_quat_w), - wp.to_torch(env.scene["factory_gear_medium"].data.root_link_quat_w), - wp.to_torch(env.scene["factory_gear_large"].data.root_link_quat_w), + env.scene["factory_gear_small"].data.root_link_quat_w.torch, + env.scene["factory_gear_medium"].data.root_link_quat_w.torch, + env.scene["factory_gear_large"].data.root_link_quat_w.torch, ], dim=1, )[env_ids] @@ -306,8 +306,8 @@ def __call__( ) # Get end effector pose - eef_pos = wp.to_torch(self.robot_asset.data.body_pos_w)[env_ids, self.eef_idx] - eef_quat = wp.to_torch(self.robot_asset.data.body_quat_w)[env_ids, self.eef_idx] + eef_pos = self.robot_asset.data.body_pos_w.torch[env_ids, self.eef_idx] + eef_quat = self.robot_asset.data.body_quat_w.torch[env_ids, self.eef_idx] # Compute pose error pos_error, axis_angle_error = fc.get_pose_error( @@ -349,7 +349,7 @@ def __call__( self.robot_asset.write_joint_velocity_to_sim_index(velocity=joint_vel, env_ids=env_ids) # Set gripper to grasp position - joint_pos = wp.to_torch(self.robot_asset.data.joint_pos)[env_ids].clone() + joint_pos = self.robot_asset.data.joint_pos.torch[env_ids].clone() # Get gear types for all environments all_gear_types = gear_type_manager.get_all_gear_types() @@ -447,8 +447,8 @@ def __call__( asset_names_to_process = [self.base_asset_name] + self.gear_asset_names for asset_name in asset_names_to_process: asset: RigidObject | Articulation = env.scene[asset_name] - default_root_pose = wp.to_torch(asset.data.default_root_pose)[env_ids].clone() - default_root_vel = wp.to_torch(asset.data.default_root_vel)[env_ids].clone() + default_root_pose = asset.data.default_root_pose.torch[env_ids].clone() + default_root_vel = asset.data.default_root_vel.torch[env_ids].clone() positions = default_root_pose[:, 0:3] + env.scene.env_origins[env_ids] + rand_pose_samples[:, 0:3] orientations = math_utils.quat_mul(default_root_pose[:, 3:7], orientations_delta) velocities = default_root_vel + rand_vel_samples diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/deploy/mdp/observations.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/deploy/mdp/observations.py index e70b77d7d8f9..744182befd15 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/deploy/mdp/observations.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/deploy/mdp/observations.py @@ -10,7 +10,6 @@ from typing import TYPE_CHECKING import torch -import warp as wp from isaaclab.managers import ManagerTermBase, ObservationTermCfg, SceneEntityCfg from isaaclab.utils.math import combine_frame_transforms @@ -128,8 +127,8 @@ def __call__( gear_type_indices = gear_type_manager.get_all_gear_type_indices() # Get base gear position and orientation - base_pos = wp.to_torch(self.asset.data.root_pos_w) - base_quat = wp.to_torch(self.asset.data.root_quat_w) + base_pos = self.asset.data.root_pos_w.torch + base_quat = self.asset.data.root_quat_w.torch # Update offsets using vectorized indexing self.offsets_buffer = self.gear_offsets_stacked[gear_type_indices] @@ -182,7 +181,7 @@ def __call__( Gear shaft orientation tensor of shape (num_envs, 4) """ # Get base quaternion - base_quat = wp.to_torch(self.asset.data.root_quat_w) + base_quat = self.asset.data.root_quat_w.torch # Ensure w component is positive (q and -q represent the same rotation) # Pick one canonical form to reduce observation variation seen by the policy @@ -250,9 +249,9 @@ def __call__(self, env: ManagerBasedRLEnv) -> torch.Tensor: # Stack all gear positions all_gear_positions = torch.stack( [ - wp.to_torch(self.gear_assets["gear_small"].data.root_pos_w), - wp.to_torch(self.gear_assets["gear_medium"].data.root_pos_w), - wp.to_torch(self.gear_assets["gear_large"].data.root_pos_w), + self.gear_assets["gear_small"].data.root_pos_w.torch, + self.gear_assets["gear_medium"].data.root_pos_w.torch, + self.gear_assets["gear_large"].data.root_pos_w.torch, ], dim=1, ) @@ -323,9 +322,9 @@ def __call__(self, env: ManagerBasedRLEnv) -> torch.Tensor: # Stack all gear quaternions all_gear_quat = torch.stack( [ - wp.to_torch(self.gear_assets["gear_small"].data.root_quat_w), - wp.to_torch(self.gear_assets["gear_medium"].data.root_quat_w), - wp.to_torch(self.gear_assets["gear_large"].data.root_quat_w), + self.gear_assets["gear_small"].data.root_quat_w.torch, + self.gear_assets["gear_medium"].data.root_quat_w.torch, + self.gear_assets["gear_large"].data.root_quat_w.torch, ], dim=1, ) diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/deploy/mdp/rewards.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/deploy/mdp/rewards.py index 2204a0186bc0..2eff6a6e0bca 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/deploy/mdp/rewards.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/deploy/mdp/rewards.py @@ -10,7 +10,6 @@ from typing import TYPE_CHECKING import torch -import warp as wp from isaaclab.managers import ManagerTermBase, RewardTermCfg, SceneEntityCfg from isaaclab.utils.math import combine_frame_transforms @@ -73,8 +72,8 @@ def __call__( des_quat_w = command[:, 3:7] # Get current pose from frame transformer - curr_pos_w = wp.to_torch(asset.data.target_pos_source)[:, 0] - curr_quat_w = wp.to_torch(asset.data.target_quat_source)[:, 0] + curr_pos_w = asset.data.target_pos_source.torch[:, 0] + curr_quat_w = asset.data.target_quat_source.torch[:, 0] # Compute keypoint distance keypoint_dist_sep = self.keypoint_computer.compute( @@ -144,8 +143,8 @@ def __call__( des_quat_w = command[:, 3:7] # Get current pose from frame transformer - curr_pos_w = wp.to_torch(asset.data.target_pos_source)[:, 0] - curr_quat_w = wp.to_torch(asset.data.target_quat_source)[:, 0] + curr_pos_w = asset.data.target_pos_source.torch[:, 0] + curr_quat_w = asset.data.target_quat_source.torch[:, 0] # Compute keypoint distance keypoint_dist_sep = self.keypoint_computer.compute( @@ -227,8 +226,8 @@ def __call__( Mean keypoint distance tensor of shape (num_envs,) """ # Get current pose of asset_1 (RigidObject) - curr_pos_1 = wp.to_torch(self.asset_1.data.body_pos_w)[:, 0] - curr_quat_1 = wp.to_torch(self.asset_1.data.body_quat_w)[:, 0] + curr_pos_1 = self.asset_1.data.body_pos_w.torch[:, 0] + curr_quat_1 = self.asset_1.data.body_quat_w.torch[:, 0] # Check if gear type manager exists if not hasattr(env, "_gear_type_manager"): @@ -244,18 +243,18 @@ def __call__( # Stack all gear positions and quaternions all_gear_pos = torch.stack( [ - wp.to_torch(self.gear_assets["gear_small"].data.body_pos_w)[:, 0], - wp.to_torch(self.gear_assets["gear_medium"].data.body_pos_w)[:, 0], - wp.to_torch(self.gear_assets["gear_large"].data.body_pos_w)[:, 0], + self.gear_assets["gear_small"].data.body_pos_w.torch[:, 0], + self.gear_assets["gear_medium"].data.body_pos_w.torch[:, 0], + self.gear_assets["gear_large"].data.body_pos_w.torch[:, 0], ], dim=1, ) all_gear_quat = torch.stack( [ - wp.to_torch(self.gear_assets["gear_small"].data.body_quat_w)[:, 0], - wp.to_torch(self.gear_assets["gear_medium"].data.body_quat_w)[:, 0], - wp.to_torch(self.gear_assets["gear_large"].data.body_quat_w)[:, 0], + self.gear_assets["gear_small"].data.body_quat_w.torch[:, 0], + self.gear_assets["gear_medium"].data.body_quat_w.torch[:, 0], + self.gear_assets["gear_large"].data.body_quat_w.torch[:, 0], ], dim=1, ) @@ -333,8 +332,8 @@ def __call__( Exponential keypoint reward tensor of shape (num_envs,) """ # Get current pose of asset_1 (RigidObject) - curr_pos_1 = wp.to_torch(self.asset_1.data.body_pos_w)[:, 0] - curr_quat_1 = wp.to_torch(self.asset_1.data.body_quat_w)[:, 0] + curr_pos_1 = self.asset_1.data.body_pos_w.torch[:, 0] + curr_quat_1 = self.asset_1.data.body_quat_w.torch[:, 0] # Check if gear type manager exists if not hasattr(env, "_gear_type_manager"): @@ -350,18 +349,18 @@ def __call__( # Stack all gear positions and quaternions all_gear_pos = torch.stack( [ - wp.to_torch(self.gear_assets["gear_small"].data.body_pos_w)[:, 0], - wp.to_torch(self.gear_assets["gear_medium"].data.body_pos_w)[:, 0], - wp.to_torch(self.gear_assets["gear_large"].data.body_pos_w)[:, 0], + self.gear_assets["gear_small"].data.body_pos_w.torch[:, 0], + self.gear_assets["gear_medium"].data.body_pos_w.torch[:, 0], + self.gear_assets["gear_large"].data.body_pos_w.torch[:, 0], ], dim=1, ) all_gear_quat = torch.stack( [ - wp.to_torch(self.gear_assets["gear_small"].data.body_quat_w)[:, 0], - wp.to_torch(self.gear_assets["gear_medium"].data.body_quat_w)[:, 0], - wp.to_torch(self.gear_assets["gear_large"].data.body_quat_w)[:, 0], + self.gear_assets["gear_small"].data.body_quat_w.torch[:, 0], + self.gear_assets["gear_medium"].data.body_quat_w.torch[:, 0], + self.gear_assets["gear_large"].data.body_quat_w.torch[:, 0], ], dim=1, ) diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/deploy/mdp/terminations.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/deploy/mdp/terminations.py index 1fd2b1ad9bc3..9f217c316372 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/deploy/mdp/terminations.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/deploy/mdp/terminations.py @@ -11,7 +11,6 @@ from typing import TYPE_CHECKING import torch -import warp as wp import isaaclab.utils.math as math_utils from isaaclab.managers import ManagerTermBase, SceneEntityCfg, TerminationTermCfg @@ -162,16 +161,16 @@ def __call__( self.gear_type_indices = gear_type_manager.get_all_gear_type_indices() # Get end effector position - eef_pos_world = wp.to_torch(self.robot_asset.data.body_link_pos_w)[:, self.eef_idx] + eef_pos_world = self.robot_asset.data.body_link_pos_w.torch[:, self.eef_idx] # Update gear positions and quaternions in buffers - self.all_gear_pos_buffer[:, 0, :] = wp.to_torch(self.gear_assets["gear_small"].data.root_link_pos_w) - self.all_gear_pos_buffer[:, 1, :] = wp.to_torch(self.gear_assets["gear_medium"].data.root_link_pos_w) - self.all_gear_pos_buffer[:, 2, :] = wp.to_torch(self.gear_assets["gear_large"].data.root_link_pos_w) + self.all_gear_pos_buffer[:, 0, :] = self.gear_assets["gear_small"].data.root_link_pos_w.torch + self.all_gear_pos_buffer[:, 1, :] = self.gear_assets["gear_medium"].data.root_link_pos_w.torch + self.all_gear_pos_buffer[:, 2, :] = self.gear_assets["gear_large"].data.root_link_pos_w.torch - self.all_gear_quat_buffer[:, 0, :] = wp.to_torch(self.gear_assets["gear_small"].data.root_link_quat_w) - self.all_gear_quat_buffer[:, 1, :] = wp.to_torch(self.gear_assets["gear_medium"].data.root_link_quat_w) - self.all_gear_quat_buffer[:, 2, :] = wp.to_torch(self.gear_assets["gear_large"].data.root_link_quat_w) + self.all_gear_quat_buffer[:, 0, :] = self.gear_assets["gear_small"].data.root_link_quat_w.torch + self.all_gear_quat_buffer[:, 1, :] = self.gear_assets["gear_medium"].data.root_link_quat_w.torch + self.all_gear_quat_buffer[:, 2, :] = self.gear_assets["gear_large"].data.root_link_quat_w.torch # Select gear data using advanced indexing gear_pos_world = self.all_gear_pos_buffer[self.env_indices, self.gear_type_indices] @@ -303,12 +302,12 @@ def __call__( yaw_threshold_rad = torch.deg2rad(torch.tensor(yaw_threshold_deg, device=env.device)) # Get end effector orientation - eef_quat_world = wp.to_torch(self.robot_asset.data.body_link_quat_w)[:, self.eef_idx] + eef_quat_world = self.robot_asset.data.body_link_quat_w.torch[:, self.eef_idx] # Update gear quaternions in buffer - self.all_gear_quat_buffer[:, 0, :] = wp.to_torch(self.gear_assets["gear_small"].data.root_link_quat_w) - self.all_gear_quat_buffer[:, 1, :] = wp.to_torch(self.gear_assets["gear_medium"].data.root_link_quat_w) - self.all_gear_quat_buffer[:, 2, :] = wp.to_torch(self.gear_assets["gear_large"].data.root_link_quat_w) + self.all_gear_quat_buffer[:, 0, :] = self.gear_assets["gear_small"].data.root_link_quat_w.torch + self.all_gear_quat_buffer[:, 1, :] = self.gear_assets["gear_medium"].data.root_link_quat_w.torch + self.all_gear_quat_buffer[:, 2, :] = self.gear_assets["gear_large"].data.root_link_quat_w.torch # Select gear data using advanced indexing gear_quat_world = self.all_gear_quat_buffer[self.env_indices, self.gear_type_indices] diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/mdp/commands/pose_commands.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/mdp/commands/pose_commands.py index 35549df614ab..d3a43dc933ac 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/mdp/commands/pose_commands.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/mdp/commands/pose_commands.py @@ -12,7 +12,6 @@ from typing import TYPE_CHECKING import torch -import warp as wp from isaaclab.managers import CommandTerm from isaaclab.utils.math import combine_frame_transforms, compute_pose_error, quat_from_euler_xyz, quat_unique @@ -108,13 +107,13 @@ def command(self) -> torch.Tensor: def _update_metrics(self): # transform command from base frame to simulation world frame self.pose_command_w[:, :3], self.pose_command_w[:, 3:] = combine_frame_transforms( - wp.to_torch(self.robot.data.root_pos_w), - wp.to_torch(self.robot.data.root_quat_w), + self.robot.data.root_pos_w.torch, + self.robot.data.root_quat_w.torch, self.pose_command_b[:, :3], self.pose_command_b[:, 3:], ) # compute the error - object_root_pose_w = wp.to_torch(self.object.data.root_link_pose_w) + object_root_pose_w = self.object.data.root_link_pose_w.torch pos_error, rot_error = compute_pose_error( self.pose_command_w[:, :3], self.pose_command_w[:, 3:], @@ -129,7 +128,7 @@ def _update_metrics(self): success_id &= self.metrics["orientation_error"] < 0.5 if self.success_vis_asset is not None: self.success_visualizer.visualize( - wp.to_torch(self.success_vis_asset.data.root_pos_w), marker_indices=success_id.int() + self.success_vis_asset.data.root_pos_w.torch, marker_indices=success_id.int() ) def _resample_command(self, env_ids: Sequence[int]): @@ -176,11 +175,11 @@ def _debug_vis_callback(self, event): # -- goal pose self.goal_visualizer.visualize(self.pose_command_w[:, :3], self.pose_command_w[:, 3:]) # -- current object pose - obj_pos = wp.to_torch(self.object.data.root_pos_w) - obj_quat = wp.to_torch(self.object.data.root_quat_w) + obj_pos = self.object.data.root_pos_w.torch + obj_quat = self.object.data.root_quat_w.torch self.curr_visualizer.visualize(obj_pos, obj_quat) else: - obj_pos = wp.to_torch(self.object.data.root_pos_w) + obj_pos = self.object.data.root_pos_w.torch distance = torch.linalg.norm(self.pose_command_w[:, :3] - obj_pos, dim=1) success_id = (distance < 0.05).int() # note: since marker indices for position is 1(far) and 2(near), we can simply shift the success_id by 1. diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/mdp/observations.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/mdp/observations.py index c3f0777f6d63..7fa204472691 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/mdp/observations.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/mdp/observations.py @@ -8,7 +8,6 @@ from typing import TYPE_CHECKING import torch -import warp as wp from isaaclab.managers import ManagerTermBase, SceneEntityCfg from isaaclab.utils.math import quat_apply, quat_apply_inverse, quat_inv, quat_mul, subtract_frame_transforms @@ -36,9 +35,7 @@ def object_pos_b( """ robot: RigidObject = env.scene[robot_cfg.name] object: RigidObject = env.scene[object_cfg.name] - return quat_apply_inverse( - wp.to_torch(robot.data.root_quat_w), wp.to_torch(object.data.root_pos_w) - wp.to_torch(robot.data.root_pos_w) - ) + return quat_apply_inverse(robot.data.root_quat_w.torch, object.data.root_pos_w.torch - robot.data.root_pos_w.torch) def object_quat_b( @@ -58,7 +55,7 @@ def object_quat_b( """ robot: RigidObject = env.scene[robot_cfg.name] object: RigidObject = env.scene[object_cfg.name] - return quat_mul(quat_inv(wp.to_torch(robot.data.root_quat_w)), wp.to_torch(object.data.root_quat_w)) + return quat_mul(quat_inv(robot.data.root_quat_w.torch), object.data.root_quat_w.torch) def body_state_b( @@ -82,18 +79,14 @@ def body_state_b( body_asset: Articulation = env.scene[body_asset_cfg.name] base_asset: Articulation = env.scene[base_asset_cfg.name] # get world pose of bodies - body_pos_w = wp.to_torch(body_asset.data.body_pos_w)[:, body_asset_cfg.body_ids].view(-1, 3) - body_quat_w = wp.to_torch(body_asset.data.body_quat_w)[:, body_asset_cfg.body_ids].view(-1, 4) - body_lin_vel_w = wp.to_torch(body_asset.data.body_lin_vel_w)[:, body_asset_cfg.body_ids].view(-1, 3) - body_ang_vel_w = wp.to_torch(body_asset.data.body_ang_vel_w)[:, body_asset_cfg.body_ids].view(-1, 3) + body_pos_w = body_asset.data.body_pos_w.torch[:, body_asset_cfg.body_ids].view(-1, 3) + body_quat_w = body_asset.data.body_quat_w.torch[:, body_asset_cfg.body_ids].view(-1, 4) + body_lin_vel_w = body_asset.data.body_lin_vel_w.torch[:, body_asset_cfg.body_ids].view(-1, 3) + body_ang_vel_w = body_asset.data.body_ang_vel_w.torch[:, body_asset_cfg.body_ids].view(-1, 3) num_bodies = int(body_pos_w.shape[0] / env.num_envs) # get world pose of base frame - root_pos_w = ( - wp.to_torch(base_asset.data.root_link_pos_w).unsqueeze(1).repeat_interleave(num_bodies, dim=1).view(-1, 3) - ) - root_quat_w = ( - wp.to_torch(base_asset.data.root_link_quat_w).unsqueeze(1).repeat_interleave(num_bodies, dim=1).view(-1, 4) - ) + root_pos_w = base_asset.data.root_link_pos_w.torch.unsqueeze(1).repeat_interleave(num_bodies, dim=1).view(-1, 3) + root_quat_w = base_asset.data.root_link_quat_w.torch.unsqueeze(1).repeat_interleave(num_bodies, dim=1).view(-1, 4) # transform from world body pose to local body pose body_pos_b, body_quat_b = subtract_frame_transforms(root_pos_w, root_quat_w, body_pos_w, body_quat_w) body_lin_vel_b = quat_apply_inverse(root_quat_w, body_lin_vel_w) @@ -170,11 +163,11 @@ def __call__( Returns: Tensor of shape ``(num_envs, num_points, 3)`` or flattened if requested. """ - ref_pos_w = wp.to_torch(self.ref_asset.data.root_pos_w).unsqueeze(1).repeat(1, num_points, 1) - ref_quat_w = wp.to_torch(self.ref_asset.data.root_quat_w).unsqueeze(1).repeat(1, num_points, 1) + ref_pos_w = self.ref_asset.data.root_pos_w.torch.unsqueeze(1).repeat(1, num_points, 1) + ref_quat_w = self.ref_asset.data.root_quat_w.torch.unsqueeze(1).repeat(1, num_points, 1) - object_pos_w = wp.to_torch(self.object.data.root_pos_w).unsqueeze(1).repeat(1, num_points, 1) - object_quat_w = wp.to_torch(self.object.data.root_quat_w).unsqueeze(1).repeat(1, num_points, 1) + object_pos_w = self.object.data.root_pos_w.torch.unsqueeze(1).repeat(1, num_points, 1) + object_quat_w = self.object.data.root_quat_w.torch.unsqueeze(1).repeat(1, num_points, 1) # apply rotation + translation self.points_w = quat_apply(object_quat_w, self.points_local) + object_pos_w if visualize: @@ -199,12 +192,10 @@ def fingers_contact_force_b( Tensor of shape ``(num_envs, 3 * num_sensors)`` with forces stacked horizontally as ``[fx, fy, fz]`` per sensor. """ - force_w = [ - wp.to_torch(env.scene.sensors[name].data.force_matrix_w).view(env.num_envs, 3) for name in contact_sensor_names - ] + force_w = [env.scene.sensors[name].data.force_matrix_w.torch.view(env.num_envs, 3) for name in contact_sensor_names] force_w = torch.stack(force_w, dim=1) robot: Articulation = env.scene[asset_cfg.name] - root_link_quat_w = wp.to_torch(robot.data.root_link_quat_w) + root_link_quat_w = robot.data.root_link_quat_w.torch forces_b = quat_apply_inverse(root_link_quat_w.unsqueeze(1).repeat(1, force_w.shape[1], 1), force_w) return forces_b.view(env.num_envs, -1) diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/mdp/rewards.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/mdp/rewards.py index 2958352ded1f..6793fc1a2ef5 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/mdp/rewards.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/mdp/rewards.py @@ -9,7 +9,6 @@ from typing import TYPE_CHECKING import torch -import warp as wp from isaaclab.managers import ManagerTermBase, SceneEntityCfg from isaaclab.utils import math as math_utils @@ -57,8 +56,8 @@ def object_ee_distance( """ asset: RigidObject = env.scene[asset_cfg.name] obj: RigidObject = env.scene[object_cfg.name] - asset_pos = wp.to_torch(asset.data.body_pos_w)[:, asset_cfg.body_ids] - object_pos = wp.to_torch(obj.data.root_pos_w) + asset_pos = asset.data.body_pos_w.torch[:, asset_cfg.body_ids] + object_pos = obj.data.root_pos_w.torch distance = torch.linalg.norm(asset_pos - object_pos[:, None, :], dim=-1).max(dim=-1).values contact_bonus = contacts(env, contact_threshold, thumb_name, finger_names).float().clamp(0.1, 1.0) return (1 - torch.tanh(distance / std)) * contact_bonus @@ -66,7 +65,7 @@ def object_ee_distance( def _contact_force_mag(sensor: ContactSensor, num_envs: int) -> torch.Tensor: """Extract per-environment contact force magnitude from a sensor's force_matrix_w.""" - force = wp.to_torch(sensor.data.force_matrix_w).view(num_envs, 3) + force = sensor.data.force_matrix_w.torch.view(num_envs, 3) return torch.linalg.norm(force, dim=-1) @@ -153,16 +152,16 @@ def __call__( obj: RigidObject = env.scene[align_asset_cfg.name] command = env.command_manager.get_command(command_name) des_pos_w, des_quat_w = combine_frame_transforms( - wp.to_torch(asset.data.root_pos_w), - wp.to_torch(asset.data.root_quat_w), + asset.data.root_pos_w.torch, + asset.data.root_quat_w.torch, command[:, :3], command[:, 3:7], ) pos_err, rot_err = compute_pose_error( des_pos_w, des_quat_w, - wp.to_torch(obj.data.root_pos_w), - wp.to_torch(obj.data.root_quat_w), + obj.data.root_pos_w.torch, + obj.data.root_quat_w.torch, ) pos_dist = torch.linalg.norm(pos_err, dim=1) contact_mask = contacts(env, contact_threshold, thumb_name, finger_names) @@ -195,11 +194,11 @@ def position_command_error_tanh( command = env.command_manager.get_command(command_name) des_pos_b = command[:, :3] des_pos_w, _ = combine_frame_transforms( - wp.to_torch(asset.data.root_pos_w), - wp.to_torch(asset.data.root_quat_w), + asset.data.root_pos_w.torch, + asset.data.root_quat_w.torch, des_pos_b, ) - distance = torch.linalg.norm(wp.to_torch(obj.data.root_pos_w) - des_pos_w, dim=1) + distance = torch.linalg.norm(obj.data.root_pos_w.torch - des_pos_w, dim=1) return (1 - torch.tanh(distance / std)) * contacts(env, contact_threshold, thumb_name, finger_names).float() @@ -219,8 +218,8 @@ def orientation_command_error_tanh( obj: RigidObject = env.scene[align_asset_cfg.name] command = env.command_manager.get_command(command_name) des_quat_b = command[:, 3:7] - root_state = wp.to_torch(asset.data.root_state_w) + root_state = asset.data.root_state_w.torch des_quat_w = math_utils.quat_mul(root_state[:, 3:7], des_quat_b) - quat_distance = math_utils.quat_error_magnitude(wp.to_torch(obj.data.root_quat_w), des_quat_w) + quat_distance = math_utils.quat_error_magnitude(obj.data.root_quat_w.torch, des_quat_w) return (1 - torch.tanh(quat_distance / std)) * contacts(env, contact_threshold, thumb_name, finger_names).float() diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/mdp/terminations.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/mdp/terminations.py index 8490847814d6..13fd1d9d3666 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/mdp/terminations.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/mdp/terminations.py @@ -14,7 +14,6 @@ from typing import TYPE_CHECKING import torch -import warp as wp from isaaclab.managers import ManagerTermBase, SceneEntityCfg, TerminationTermCfg @@ -51,7 +50,7 @@ def __call__( asset_cfg: SceneEntityCfg = SceneEntityCfg("object"), in_bound_range: dict[str, tuple[float, float]] = {}, ) -> torch.Tensor: - pos_w = wp.to_torch(self._object.data.root_pos_w) + pos_w = self._object.data.root_pos_w.torch return ((pos_w < self._lower) | (pos_w > self._upper)).any(dim=1) @@ -59,6 +58,6 @@ def abnormal_robot_state(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg = Sce """Terminating environment when violation of velocity limits detects, this usually indicates unstable physics caused by very bad, or aggressive action""" robot: Articulation = env.scene[asset_cfg.name] - joint_vel = wp.to_torch(robot.data.joint_vel) - joint_vel_limits = wp.to_torch(robot.data.joint_vel_limits) + joint_vel = robot.data.joint_vel.torch + joint_vel_limits = robot.data.joint_vel_limits.torch return (joint_vel.abs() > (joint_vel_limits * 2)).any(dim=1) diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/inhand/mdp/commands/orientation_command.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/inhand/mdp/commands/orientation_command.py index 468ae2764743..7a2bf70b64ee 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/inhand/mdp/commands/orientation_command.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/inhand/mdp/commands/orientation_command.py @@ -11,7 +11,6 @@ from typing import TYPE_CHECKING import torch -import warp as wp import isaaclab.utils.math as math_utils from isaaclab.managers import CommandTerm @@ -59,7 +58,7 @@ def __init__(self, cfg: InHandReOrientationCommandCfg, env: ManagerBasedRLEnv): # create buffers to store the command # -- command: (x, y, z) init_pos_offset = torch.tensor(cfg.init_pos_offset, dtype=torch.float, device=self.device) - self.pos_command_e = wp.to_torch(self.object.data.default_root_pose)[:, :3] + init_pos_offset + self.pos_command_e = self.object.data.default_root_pose.torch[:, :3] + init_pos_offset self.pos_command_w = self.pos_command_e + self._env.scene.env_origins # -- orientation: (x, y, z, w) self.quat_command_w = torch.zeros(self.num_envs, 4, device=self.device) @@ -97,11 +96,11 @@ def _update_metrics(self): # logs data # -- compute the orientation error self.metrics["orientation_error"] = math_utils.quat_error_magnitude( - wp.to_torch(self.object.data.root_quat_w), self.quat_command_w + self.object.data.root_quat_w.torch, self.quat_command_w ) # -- compute the position error self.metrics["position_error"] = torch.linalg.norm( - wp.to_torch(self.object.data.root_pos_w) - self.pos_command_w, dim=1 + self.object.data.root_pos_w.torch - self.pos_command_w, dim=1 ) # -- compute the number of consecutive successes successes = self.metrics["orientation_error"] < self.cfg.orientation_success_threshold diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/inhand/mdp/events.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/inhand/mdp/events.py index 4382e1daa08e..c118cd16a506 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/inhand/mdp/events.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/inhand/mdp/events.py @@ -10,7 +10,6 @@ from typing import TYPE_CHECKING, Literal import torch -import warp as wp from isaaclab.managers import EventTermCfg, ManagerTermBase, SceneEntityCfg from isaaclab.utils.math import sample_uniform @@ -78,11 +77,11 @@ def __init__(self, cfg: EventTermCfg, env: ManagerBasedEnv): # extract the used quantities (to enable type-hinting) self._asset: Articulation = env.scene[asset_cfg.name] - default_joint_pos = wp.to_torch(self._asset.data.default_joint_pos)[0] - default_joint_vel = wp.to_torch(self._asset.data.default_joint_vel)[0] + default_joint_pos = self._asset.data.default_joint_pos.torch[0] + default_joint_vel = self._asset.data.default_joint_vel.torch[0] # create buffers to store the joint position range - self._pos_ranges = wp.to_torch(self._asset.data.soft_joint_pos_limits)[0].clone() + self._pos_ranges = self._asset.data.soft_joint_pos_limits.torch[0].clone() # parse joint position ranges pos_joint_ids = [] for joint_name, joint_range in cfg.params["position_range"].items(): @@ -114,7 +113,7 @@ def __init__(self, cfg: EventTermCfg, env: ManagerBasedEnv): self._pos_ranges = self._pos_ranges[self._pos_joint_ids] # create buffers to store the joint velocity range - soft_joint_vel_limits_torch = wp.to_torch(self._asset.data.soft_joint_vel_limits)[0] + soft_joint_vel_limits_torch = self._asset.data.soft_joint_vel_limits.torch[0] self._vel_ranges = torch.stack([-soft_joint_vel_limits_torch, soft_joint_vel_limits_torch], dim=1) # parse joint velocity ranges vel_joint_ids = [] @@ -157,8 +156,8 @@ def __call__( operation: Literal["abs", "scale"] = "abs", ): # get default joint state - joint_pos = wp.to_torch(self._asset.data.default_joint_pos)[env_ids].clone() - joint_vel = wp.to_torch(self._asset.data.default_joint_vel)[env_ids].clone() + joint_pos = self._asset.data.default_joint_pos.torch[env_ids].clone() + joint_vel = self._asset.data.default_joint_vel.torch[env_ids].clone() # sample random joint positions for each joint if len(self._pos_joint_ids) > 0: @@ -167,7 +166,7 @@ def __call__( self._pos_ranges[:, 0], self._pos_ranges[:, 1], joint_pos_shape, device=joint_pos.device ) # clip the joint positions to the joint limits - joint_pos_limits = wp.to_torch(self._asset.data.soft_joint_pos_limits)[0, self._pos_joint_ids] + joint_pos_limits = self._asset.data.soft_joint_pos_limits.torch[0, self._pos_joint_ids] joint_pos = joint_pos.clamp(joint_pos_limits[:, 0], joint_pos_limits[:, 1]) # sample random joint velocities for each joint @@ -177,7 +176,7 @@ def __call__( self._vel_ranges[:, 0], self._vel_ranges[:, 1], joint_vel_shape, device=joint_vel.device ) # clip the joint velocities to the joint limits - joint_vel_limits = wp.to_torch(self._asset.data.soft_joint_vel_limits)[0, self._vel_joint_ids] + joint_vel_limits = self._asset.data.soft_joint_vel_limits.torch[0, self._vel_joint_ids] joint_vel = joint_vel.clamp(-joint_vel_limits, joint_vel_limits) # set into the physics simulation diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/inhand/mdp/observations.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/inhand/mdp/observations.py index 61f6d2543d39..57620602b6fd 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/inhand/mdp/observations.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/inhand/mdp/observations.py @@ -10,7 +10,6 @@ from typing import TYPE_CHECKING import torch -import warp as wp import isaaclab.utils.math as math_utils from isaaclab.managers import SceneEntityCfg @@ -35,7 +34,7 @@ def goal_quat_diff( # obtain the orientations goal_quat_w = command_term.command[:, 3:7] - asset_quat_w = wp.to_torch(asset.data.root_quat_w) + asset_quat_w = asset.data.root_quat_w.torch # compute quaternion difference quat = math_utils.quat_mul(asset_quat_w, math_utils.quat_conjugate(goal_quat_w)) diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/inhand/mdp/rewards.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/inhand/mdp/rewards.py index da02a20c5aa3..f89b2f613b51 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/inhand/mdp/rewards.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/inhand/mdp/rewards.py @@ -10,7 +10,6 @@ from typing import TYPE_CHECKING import torch -import warp as wp import isaaclab.utils.math as math_utils from isaaclab.managers import SceneEntityCfg @@ -44,7 +43,7 @@ def success_bonus( # obtain the threshold for the orientation error threshold = command_term.cfg.orientation_success_threshold # calculate the orientation error - dtheta = math_utils.quat_error_magnitude(wp.to_torch(asset.data.root_quat_w), goal_quat_w) + dtheta = math_utils.quat_error_magnitude(asset.data.root_quat_w.torch, goal_quat_w) return dtheta <= threshold @@ -68,7 +67,7 @@ def track_pos_l2( # obtain the goal position goal_pos_e = command_term.command[:, 0:3] # obtain the object position in the environment frame - object_pos_e = wp.to_torch(asset.data.root_pos_w) - env.scene.env_origins + object_pos_e = asset.data.root_pos_w.torch - env.scene.env_origins return torch.linalg.norm(goal_pos_e - object_pos_e, ord=2, dim=-1) @@ -96,6 +95,6 @@ def track_orientation_inv_l2( # obtain the goal orientation goal_quat_w = command_term.command[:, 3:7] # calculate the orientation error - dtheta = math_utils.quat_error_magnitude(wp.to_torch(asset.data.root_quat_w), goal_quat_w) + dtheta = math_utils.quat_error_magnitude(asset.data.root_quat_w.torch, goal_quat_w) return 1.0 / (dtheta + rot_eps) diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/inhand/mdp/terminations.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/inhand/mdp/terminations.py index 895e09581bd3..c2c7a02f6797 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/inhand/mdp/terminations.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/inhand/mdp/terminations.py @@ -10,7 +10,6 @@ from typing import TYPE_CHECKING import torch -import warp as wp from isaaclab.managers import SceneEntityCfg @@ -55,7 +54,7 @@ def object_away_from_goal( asset = env.scene[object_cfg.name] # object pos - asset_pos_e = wp.to_torch(asset.data.root_pos_w) - env.scene.env_origins + asset_pos_e = asset.data.root_pos_w.torch - env.scene.env_origins goal_pos_e = command_term.command[:, :3] return torch.linalg.norm(asset_pos_e - goal_pos_e, ord=2, dim=1) > threshold @@ -83,6 +82,6 @@ def object_away_from_robot( object = env.scene[object_cfg.name] # compute distance - dist = torch.linalg.norm(wp.to_torch(robot.data.root_pos_w) - wp.to_torch(object.data.root_pos_w), dim=1) + dist = torch.linalg.norm(robot.data.root_pos_w.torch - object.data.root_pos_w.torch, dim=1) return dist > threshold diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/lift/mdp/observations.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/lift/mdp/observations.py index 1ccf60969814..06659683d770 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/lift/mdp/observations.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/lift/mdp/observations.py @@ -8,7 +8,6 @@ from typing import TYPE_CHECKING import torch -import warp as wp from isaaclab.managers import SceneEntityCfg from isaaclab.utils.math import subtract_frame_transforms @@ -26,8 +25,6 @@ def object_position_in_robot_root_frame( """The position of the object in the robot's root frame.""" robot: RigidObject = env.scene[robot_cfg.name] object: RigidObject = env.scene[object_cfg.name] - object_pos_w = wp.to_torch(object.data.root_pos_w)[:, :3] - object_pos_b, _ = subtract_frame_transforms( - wp.to_torch(robot.data.root_pos_w), wp.to_torch(robot.data.root_quat_w), object_pos_w - ) + object_pos_w = object.data.root_pos_w.torch[:, :3] + object_pos_b, _ = subtract_frame_transforms(robot.data.root_pos_w.torch, robot.data.root_quat_w.torch, object_pos_w) return object_pos_b diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/lift/mdp/rewards.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/lift/mdp/rewards.py index 2f1ccaf0287e..33ab3e7f32be 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/lift/mdp/rewards.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/lift/mdp/rewards.py @@ -8,7 +8,6 @@ from typing import TYPE_CHECKING import torch -import warp as wp from isaaclab.managers import SceneEntityCfg from isaaclab.utils.math import combine_frame_transforms @@ -24,7 +23,7 @@ def object_is_lifted( ) -> torch.Tensor: """Reward the agent for lifting the object above the minimal height.""" object: RigidObject = env.scene[object_cfg.name] - return torch.where(wp.to_torch(object.data.root_pos_w)[:, 2] > minimal_height, 1.0, 0.0) + return torch.where(object.data.root_pos_w.torch[:, 2] > minimal_height, 1.0, 0.0) def object_ee_distance( @@ -38,9 +37,9 @@ def object_ee_distance( object: RigidObject = env.scene[object_cfg.name] ee_frame: FrameTransformer = env.scene[ee_frame_cfg.name] # Target object position: (num_envs, 3) - cube_pos_w = wp.to_torch(object.data.root_pos_w) + cube_pos_w = object.data.root_pos_w.torch # End-effector position: (num_envs, 3) - ee_w = wp.to_torch(ee_frame.data.target_pos_w)[..., 0, :] + ee_w = ee_frame.data.target_pos_w.torch[..., 0, :] # Distance of the end-effector to the object: (num_envs,) object_ee_distance = torch.linalg.norm(cube_pos_w - ee_w, dim=1) @@ -62,11 +61,9 @@ def object_goal_distance( command = env.command_manager.get_command(command_name) # compute the desired position in the world frame des_pos_b = command[:, :3] - des_pos_w, _ = combine_frame_transforms( - wp.to_torch(robot.data.root_pos_w), wp.to_torch(robot.data.root_quat_w), des_pos_b - ) + des_pos_w, _ = combine_frame_transforms(robot.data.root_pos_w.torch, robot.data.root_quat_w.torch, des_pos_b) # distance of the end-effector to the object: (num_envs,) - object_pos_w = wp.to_torch(object.data.root_pos_w) + object_pos_w = object.data.root_pos_w.torch distance = torch.linalg.norm(des_pos_w - object_pos_w, dim=1) # rewarded if the object is lifted above the threshold return (object_pos_w[:, 2] > minimal_height) * (1 - torch.tanh(distance / std)) diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/lift/mdp/terminations.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/lift/mdp/terminations.py index 0f1a392fe5a0..5342ab669758 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/lift/mdp/terminations.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/lift/mdp/terminations.py @@ -14,7 +14,6 @@ from typing import TYPE_CHECKING import torch -import warp as wp from isaaclab.managers import SceneEntityCfg from isaaclab.utils.math import combine_frame_transforms @@ -48,11 +47,11 @@ def object_reached_goal( # compute the desired position in the world frame des_pos_b = command[:, :3] # Convert to torch for combine_frame_transforms (robot data may be Warp arrays under Newton) - root_pos_w = wp.to_torch(robot.data.root_pos_w) - root_quat_w = wp.to_torch(robot.data.root_quat_w) + root_pos_w = robot.data.root_pos_w.torch + root_quat_w = robot.data.root_quat_w.torch des_pos_w, _ = combine_frame_transforms(root_pos_w, root_quat_w, des_pos_b) # distance of the end-effector to the object: (num_envs,) - object_pos_w = wp.to_torch(object.data.root_pos_w) + object_pos_w = object.data.root_pos_w.torch distance = torch.linalg.norm(des_pos_w - object_pos_w[:, :3], dim=1) # rewarded if the object is lifted above the threshold diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/pick_place/mdp/observations.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/pick_place/mdp/observations.py index 87c699bcec56..c3e030b33ec0 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/pick_place/mdp/observations.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/pick_place/mdp/observations.py @@ -8,7 +8,6 @@ from typing import TYPE_CHECKING import torch -import warp as wp if TYPE_CHECKING: from isaaclab.envs import ManagerBasedRLEnv @@ -27,14 +26,14 @@ def object_obs( right_eef_to object, """ - body_pos_w = wp.to_torch(env.scene["robot"].data.body_pos_w) + body_pos_w = env.scene["robot"].data.body_pos_w.torch left_eef_idx = env.scene["robot"].data.body_names.index(left_eef_link_name) right_eef_idx = env.scene["robot"].data.body_names.index(right_eef_link_name) left_eef_pos = body_pos_w[:, left_eef_idx] - env.scene.env_origins right_eef_pos = body_pos_w[:, right_eef_idx] - env.scene.env_origins - object_pos = wp.to_torch(env.scene["object"].data.root_pos_w) - env.scene.env_origins - object_quat = wp.to_torch(env.scene["object"].data.root_quat_w) + object_pos = env.scene["object"].data.root_pos_w.torch - env.scene.env_origins + object_quat = env.scene["object"].data.root_quat_w.torch left_eef_to_object = object_pos - left_eef_pos right_eef_to_object = object_pos - right_eef_pos @@ -51,7 +50,7 @@ def object_obs( def get_eef_pos(env: ManagerBasedRLEnv, link_name: str) -> torch.Tensor: - body_pos_w = wp.to_torch(env.scene["robot"].data.body_pos_w) + body_pos_w = env.scene["robot"].data.body_pos_w.torch left_eef_idx = env.scene["robot"].data.body_names.index(link_name) left_eef_pos = body_pos_w[:, left_eef_idx] - env.scene.env_origins @@ -59,7 +58,7 @@ def get_eef_pos(env: ManagerBasedRLEnv, link_name: str) -> torch.Tensor: def get_eef_quat(env: ManagerBasedRLEnv, link_name: str) -> torch.Tensor: - body_quat_w = wp.to_torch(env.scene["robot"].data.body_quat_w) + body_quat_w = env.scene["robot"].data.body_quat_w.torch left_eef_idx = env.scene["robot"].data.body_names.index(link_name) left_eef_quat = body_quat_w[:, left_eef_idx] @@ -73,7 +72,7 @@ def get_robot_joint_state( # hand_joint_names is a list of regex, use find_joints indexes, _ = env.scene["robot"].find_joints(joint_names) indexes = torch.tensor(indexes, dtype=torch.long) - robot_joint_states = wp.to_torch(env.scene["robot"].data.joint_pos)[:, indexes] + robot_joint_states = env.scene["robot"].data.joint_pos.torch[:, indexes] return robot_joint_states @@ -81,7 +80,7 @@ def get_robot_joint_state( def get_all_robot_link_state( env: ManagerBasedRLEnv, ) -> torch.Tensor: - body_pos_w = wp.to_torch(env.scene["robot"].data.body_link_state_w)[:, :, :] + body_pos_w = env.scene["robot"].data.body_link_state_w.torch[:, :, :] all_robot_link_pos = body_pos_w return all_robot_link_pos diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/pick_place/mdp/pick_place_events.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/pick_place/mdp/pick_place_events.py index e091087bc8af..f1b4bd15de97 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/pick_place/mdp/pick_place_events.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/pick_place/mdp/pick_place_events.py @@ -8,7 +8,6 @@ from typing import TYPE_CHECKING import torch -import warp as wp import isaaclab.utils.math as math_utils from isaaclab.managers import SceneEntityCfg @@ -45,10 +44,10 @@ def reset_object_poses_nut_pour( sorting_scale = env.scene[sorting_scale_cfg.name] # get default root state - sorting_beaker_root_poses = wp.to_torch(sorting_beaker.data.default_root_pose)[env_ids].clone() - factory_nut_root_poses = wp.to_torch(factory_nut.data.default_root_pose)[env_ids].clone() - sorting_bowl_root_poses = wp.to_torch(sorting_bowl.data.default_root_pose)[env_ids].clone() - sorting_scale_root_poses = wp.to_torch(sorting_scale.data.default_root_pose)[env_ids].clone() + sorting_beaker_root_poses = sorting_beaker.data.default_root_pose.torch[env_ids].clone() + factory_nut_root_poses = factory_nut.data.default_root_pose.torch[env_ids].clone() + sorting_bowl_root_poses = sorting_bowl.data.default_root_pose.torch[env_ids].clone() + sorting_scale_root_poses = sorting_scale.data.default_root_pose.torch[env_ids].clone() # get pose ranges range_list = [pose_range.get(key, (0.0, 0.0)) for key in ["x", "y", "z", "roll", "pitch", "yaw"]] diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/pick_place/mdp/terminations.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/pick_place/mdp/terminations.py index 6c1f9161042d..7ff4732f2640 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/pick_place/mdp/terminations.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/pick_place/mdp/terminations.py @@ -14,7 +14,6 @@ from typing import TYPE_CHECKING import torch -import warp as wp from isaaclab.managers import SceneEntityCfg @@ -64,13 +63,13 @@ def task_done_pick_place( object: RigidObject = env.scene[object_cfg.name] # Extract wheel position relative to environment origin - object_x = wp.to_torch(object.data.root_pos_w)[:, 0] - env.scene.env_origins[:, 0] - object_y = wp.to_torch(object.data.root_pos_w)[:, 1] - env.scene.env_origins[:, 1] - object_height = wp.to_torch(object.data.root_pos_w)[:, 2] - env.scene.env_origins[:, 2] - object_vel = torch.abs(wp.to_torch(object.data.root_vel_w)) + object_x = object.data.root_pos_w.torch[:, 0] - env.scene.env_origins[:, 0] + object_y = object.data.root_pos_w.torch[:, 1] - env.scene.env_origins[:, 1] + object_height = object.data.root_pos_w.torch[:, 2] - env.scene.env_origins[:, 2] + object_vel = torch.abs(object.data.root_vel_w.torch) # Get right wrist position relative to environment origin - robot_body_pos_w = wp.to_torch(env.scene["robot"].data.body_pos_w) + robot_body_pos_w = env.scene["robot"].data.body_pos_w.torch right_eef_idx = env.scene["robot"].data.body_names.index(task_link_name) right_wrist_x = robot_body_pos_w[:, right_eef_idx, 0] - env.scene.env_origins[:, 0] @@ -140,11 +139,11 @@ def task_done_nut_pour( sorting_bin: RigidObject = env.scene[sorting_bin_cfg.name] # Get positions relative to environment origin - scale_pos = wp.to_torch(sorting_scale.data.root_pos_w) - env.scene.env_origins - bowl_pos = wp.to_torch(sorting_bowl.data.root_pos_w) - env.scene.env_origins - sorting_beaker_pos = wp.to_torch(sorting_beaker.data.root_pos_w) - env.scene.env_origins - nut_pos = wp.to_torch(factory_nut.data.root_pos_w) - env.scene.env_origins - bin_pos = wp.to_torch(sorting_bin.data.root_pos_w) - env.scene.env_origins + scale_pos = sorting_scale.data.root_pos_w.torch - env.scene.env_origins + bowl_pos = sorting_bowl.data.root_pos_w.torch - env.scene.env_origins + sorting_beaker_pos = sorting_beaker.data.root_pos_w.torch - env.scene.env_origins + nut_pos = factory_nut.data.root_pos_w.torch - env.scene.env_origins + bin_pos = sorting_bin.data.root_pos_w.torch - env.scene.env_origins # nut to bowl nut_to_bowl_x = torch.abs(nut_pos[:, 0] - bowl_pos[:, 0]) @@ -207,8 +206,8 @@ def task_done_exhaust_pipe( blue_sorting_bin: RigidObject = env.scene[blue_sorting_bin_cfg.name] # Get positions relative to environment origin - blue_exhaust_pipe_pos = wp.to_torch(blue_exhaust_pipe.data.root_pos_w) - env.scene.env_origins - blue_sorting_bin_pos = wp.to_torch(blue_sorting_bin.data.root_pos_w) - env.scene.env_origins + blue_exhaust_pipe_pos = blue_exhaust_pipe.data.root_pos_w.torch - env.scene.env_origins + blue_sorting_bin_pos = blue_sorting_bin.data.root_pos_w.torch - env.scene.env_origins # blue exhaust to bin blue_exhaust_to_bin_x = torch.abs(blue_exhaust_pipe_pos[:, 0] - blue_sorting_bin_pos[:, 0]) diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/place/mdp/observations.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/place/mdp/observations.py index 70c4e9391ad1..0c144db467eb 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/place/mdp/observations.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/place/mdp/observations.py @@ -28,13 +28,13 @@ def object_poses_in_base_frame( """The pose of the object in the robot base frame.""" object: RigidObject = env.scene[object_cfg.name] - pos_object_world = wp.to_torch(object.data.root_pos_w) - quat_object_world = wp.to_torch(object.data.root_quat_w) + pos_object_world = object.data.root_pos_w.torch + quat_object_world = object.data.root_quat_w.torch """The position of the robot in the world frame.""" robot: Articulation = env.scene[robot_cfg.name] - root_pos_w = wp.to_torch(robot.data.root_pos_w) - root_quat_w = wp.to_torch(robot.data.root_quat_w) + root_pos_w = robot.data.root_pos_w.torch + root_quat_w = robot.data.root_quat_w.torch pos_object_base, quat_object_base = math_utils.subtract_frame_transforms( root_pos_w, root_quat_w, pos_object_world, quat_object_world @@ -65,14 +65,12 @@ def object_grasped( ee_frame: FrameTransformer = env.scene[ee_frame_cfg.name] object: RigidObject = env.scene[object_cfg.name] - object_pos = wp.to_torch(object.data.root_pos_w) - end_effector_pos = wp.to_torch(ee_frame.data.target_pos_w)[:, 0, :] + object_pos = object.data.root_pos_w.torch + end_effector_pos = ee_frame.data.target_pos_w.torch[:, 0, :] pose_diff = torch.linalg.vector_norm(object_pos - end_effector_pos, dim=1) if "contact_grasp" in env.scene.keys() and env.scene["contact_grasp"] is not None: - contact_force_grasp = wp.to_torch( - env.scene["contact_grasp"].data.net_forces_w - ) # shape:(N, 2, 3) for two fingers + contact_force_grasp = env.scene["contact_grasp"].data.net_forces_w.torch # shape:(N, 2, 3) for two fingers contact_force_norm = torch.linalg.vector_norm( contact_force_grasp, dim=2 ) # shape:(N, 2) - force magnitude per finger @@ -84,9 +82,9 @@ def object_grasped( f"contact_grasp_{object_cfg.name}" in env.scene.keys() and env.scene[f"contact_grasp_{object_cfg.name}"] is not None ): - contact_force_object = wp.to_torch( - env.scene[f"contact_grasp_{object_cfg.name}"].data.net_forces_w - ) # shape:(N, 2, 3) for two fingers + contact_force_object = env.scene[ + f"contact_grasp_{object_cfg.name}" + ].data.net_forces_w.torch # shape:(N, 2, 3) for two fingers contact_force_norm = torch.linalg.vector_norm( contact_force_object, dim=2 ) # shape:(N, 2) - force magnitude per finger @@ -108,16 +106,12 @@ def object_grasped( gripper_joint_ids, _ = robot.find_joints(env.cfg.gripper_joint_names) grasped = torch.logical_and( grasped, - torch.abs( - torch.abs(wp.to_torch(robot.data.joint_pos)[:, gripper_joint_ids[0]]) - env.cfg.gripper_open_val - ) + torch.abs(torch.abs(robot.data.joint_pos.torch[:, gripper_joint_ids[0]]) - env.cfg.gripper_open_val) > env.cfg.gripper_threshold, ) grasped = torch.logical_and( grasped, - torch.abs( - torch.abs(wp.to_torch(robot.data.joint_pos)[:, gripper_joint_ids[1]]) - env.cfg.gripper_open_val - ) + torch.abs(torch.abs(robot.data.joint_pos.torch[:, gripper_joint_ids[1]]) - env.cfg.gripper_open_val) > env.cfg.gripper_threshold, ) else: diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/place/mdp/terminations.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/place/mdp/terminations.py index 139617d3050f..11368e935181 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/place/mdp/terminations.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/place/mdp/terminations.py @@ -37,9 +37,7 @@ def object_placed_upright( object: RigidObject = env.scene[object_cfg.name] # Compute mug euler angles of X, Y axis, to check if it is placed upright - object_euler_x, object_euler_y, _ = math_utils.euler_xyz_from_quat( - wp.to_torch(object.data.root_quat_w) - ) # (N,4) [0, 2*pi] + object_euler_x, object_euler_y, _ = math_utils.euler_xyz_from_quat(object.data.root_quat_w.torch) # (N,4) [0, 2*pi] object_euler_x_err = torch.abs(math_utils.wrap_to_pi(object_euler_x)) # (N,) object_euler_y_err = torch.abs(math_utils.wrap_to_pi(object_euler_y)) # (N,) @@ -47,7 +45,7 @@ def object_placed_upright( success = torch.logical_and(object_euler_x_err < euler_xy_threshold, object_euler_y_err < euler_xy_threshold) # Check if current mug height is greater than target height - height_success = wp.to_torch(object.data.root_pos_w)[:, 2] > target_height + height_success = object.data.root_pos_w.torch[:, 2] > target_height success = torch.logical_and(height_success, success) @@ -62,16 +60,12 @@ def object_placed_upright( gripper_joint_ids, _ = robot.find_joints(env.cfg.gripper_joint_names) success = torch.logical_and( success, - torch.abs( - torch.abs(wp.to_torch(robot.data.joint_pos)[:, gripper_joint_ids[0]]) - env.cfg.gripper_open_val - ) + torch.abs(torch.abs(robot.data.joint_pos.torch[:, gripper_joint_ids[0]]) - env.cfg.gripper_open_val) < env.cfg.gripper_threshold, ) success = torch.logical_and( success, - torch.abs( - torch.abs(wp.to_torch(robot.data.joint_pos)[:, gripper_joint_ids[1]]) - env.cfg.gripper_open_val - ) + torch.abs(torch.abs(robot.data.joint_pos.torch[:, gripper_joint_ids[1]]) - env.cfg.gripper_open_val) < env.cfg.gripper_threshold, ) else: @@ -96,7 +90,7 @@ def object_a_is_into_b( object_b: RigidObject = env.scene[object_b_cfg.name] # check object a is into object b - pos_diff = wp.to_torch(object_a.data.root_pos_w) - wp.to_torch(object_b.data.root_pos_w) + pos_diff = object_a.data.root_pos_w.torch - object_b.data.root_pos_w.torch height_dist = torch.linalg.vector_norm(pos_diff[:, 2:], dim=1) xy_dist = torch.linalg.vector_norm(pos_diff[:, :2], dim=1) @@ -116,16 +110,12 @@ def object_a_is_into_b( success = torch.logical_and( success, - torch.abs( - torch.abs(wp.to_torch(robot.data.joint_pos)[:, gripper_joint_ids[0]]) - env.cfg.gripper_open_val - ) + torch.abs(torch.abs(robot.data.joint_pos.torch[:, gripper_joint_ids[0]]) - env.cfg.gripper_open_val) < env.cfg.gripper_threshold, ) success = torch.logical_and( success, - torch.abs( - torch.abs(wp.to_torch(robot.data.joint_pos)[:, gripper_joint_ids[1]]) - env.cfg.gripper_open_val - ) + torch.abs(torch.abs(robot.data.joint_pos.torch[:, gripper_joint_ids[1]]) - env.cfg.gripper_open_val) < env.cfg.gripper_threshold, ) else: diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/reach/mdp/rewards.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/reach/mdp/rewards.py index 6ba3176b535c..3238523134d5 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/reach/mdp/rewards.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/reach/mdp/rewards.py @@ -8,7 +8,6 @@ from typing import TYPE_CHECKING import torch -import warp as wp from isaaclab.managers import SceneEntityCfg from isaaclab.utils.math import combine_frame_transforms, quat_error_magnitude, quat_mul @@ -30,10 +29,8 @@ def position_command_error(env: ManagerBasedRLEnv, command_name: str, asset_cfg: command = env.command_manager.get_command(command_name) # obtain the desired and current positions des_pos_b = command[:, :3] - des_pos_w, _ = combine_frame_transforms( - wp.to_torch(asset.data.root_pos_w), wp.to_torch(asset.data.root_quat_w), des_pos_b - ) - curr_pos_w = wp.to_torch(asset.data.body_pos_w)[:, asset_cfg.body_ids[0]] # type: ignore + des_pos_w, _ = combine_frame_transforms(asset.data.root_pos_w.torch, asset.data.root_quat_w.torch, des_pos_b) + curr_pos_w = asset.data.body_pos_w.torch[:, asset_cfg.body_ids[0]] # type: ignore return torch.linalg.norm(curr_pos_w - des_pos_w, dim=1) @@ -50,10 +47,8 @@ def position_command_error_tanh( command = env.command_manager.get_command(command_name) # obtain the desired and current positions des_pos_b = command[:, :3] - des_pos_w, _ = combine_frame_transforms( - wp.to_torch(asset.data.root_pos_w), wp.to_torch(asset.data.root_quat_w), des_pos_b - ) - curr_pos_w = wp.to_torch(asset.data.body_pos_w)[:, asset_cfg.body_ids[0]] # type: ignore + des_pos_w, _ = combine_frame_transforms(asset.data.root_pos_w.torch, asset.data.root_quat_w.torch, des_pos_b) + curr_pos_w = asset.data.body_pos_w.torch[:, asset_cfg.body_ids[0]] # type: ignore distance = torch.linalg.norm(curr_pos_w - des_pos_w, dim=1) return 1 - torch.tanh(distance / std) @@ -70,6 +65,6 @@ def orientation_command_error(env: ManagerBasedRLEnv, command_name: str, asset_c command = env.command_manager.get_command(command_name) # obtain the desired and current orientations des_quat_b = command[:, 3:7] - des_quat_w = quat_mul(wp.to_torch(asset.data.root_quat_w), des_quat_b) - curr_quat_w = wp.to_torch(asset.data.body_quat_w)[:, asset_cfg.body_ids[0]] # type: ignore + des_quat_w = quat_mul(asset.data.root_quat_w.torch, des_quat_b) + curr_quat_w = asset.data.body_quat_w.torch[:, asset_cfg.body_ids[0]] # type: ignore return quat_error_magnitude(curr_quat_w, des_quat_w) diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/mdp/franka_stack_events.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/mdp/franka_stack_events.py index 9232656eb2a6..4d148d6ba216 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/mdp/franka_stack_events.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/mdp/franka_stack_events.py @@ -58,16 +58,16 @@ def randomize_joint_by_gaussian_offset( asset: Articulation = env.scene[asset_cfg.name] # Add gaussian noise to joint states - joint_pos = wp.to_torch(asset.data.default_joint_pos)[env_ids].clone() - joint_vel = wp.to_torch(asset.data.default_joint_vel)[env_ids].clone() + joint_pos = asset.data.default_joint_pos.torch[env_ids].clone() + joint_vel = asset.data.default_joint_vel.torch[env_ids].clone() joint_pos += math_utils.sample_gaussian(mean, std, joint_pos.shape, joint_pos.device) # Clamp joint pos to limits - joint_pos_limits = wp.to_torch(asset.data.soft_joint_pos_limits)[env_ids] + joint_pos_limits = asset.data.soft_joint_pos_limits.torch[env_ids] joint_pos = joint_pos.clamp_(joint_pos_limits[..., 0], joint_pos_limits[..., 1]) # Don't noise the gripper poses - joint_pos[:, -2:] = wp.to_torch(asset.data.default_joint_pos)[env_ids, -2:] + joint_pos[:, -2:] = asset.data.default_joint_pos.torch[env_ids, -2:] # Set into the physics simulation asset.set_joint_position_target_index(target=joint_pos, env_ids=env_ids) diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/mdp/observations.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/mdp/observations.py index 2378276602d5..1269ef260ff5 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/mdp/observations.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/mdp/observations.py @@ -31,7 +31,7 @@ def cube_positions_in_world_frame( cube_3: RigidObject = env.scene[cube_3_cfg.name] return torch.cat( - (wp.to_torch(cube_1.data.root_pos_w), wp.to_torch(cube_2.data.root_pos_w), wp.to_torch(cube_3.data.root_pos_w)), + (cube_1.data.root_pos_w.torch, cube_2.data.root_pos_w.torch, cube_3.data.root_pos_w.torch), dim=1, ) @@ -54,9 +54,9 @@ def instance_randomize_cube_positions_in_world_frame( cube_2_pos_w = [] cube_3_pos_w = [] for env_id in range(env.num_envs): - cube_1_pos_w.append(wp.to_torch(cube_1.data.body_link_pos_w)[env_id, env.rigid_objects_in_focus[env_id][0], :3]) - cube_2_pos_w.append(wp.to_torch(cube_2.data.body_link_pos_w)[env_id, env.rigid_objects_in_focus[env_id][1], :3]) - cube_3_pos_w.append(wp.to_torch(cube_3.data.body_link_pos_w)[env_id, env.rigid_objects_in_focus[env_id][2], :3]) + cube_1_pos_w.append(cube_1.data.body_link_pos_w.torch[env_id, env.rigid_objects_in_focus[env_id][0], :3]) + cube_2_pos_w.append(cube_2.data.body_link_pos_w.torch[env_id, env.rigid_objects_in_focus[env_id][1], :3]) + cube_3_pos_w.append(cube_3.data.body_link_pos_w.torch[env_id, env.rigid_objects_in_focus[env_id][2], :3]) cube_1_pos_w = torch.stack(cube_1_pos_w) cube_2_pos_w = torch.stack(cube_2_pos_w) cube_3_pos_w = torch.stack(cube_3_pos_w) @@ -77,9 +77,9 @@ def cube_orientations_in_world_frame( return torch.cat( ( - wp.to_torch(cube_1.data.root_quat_w), - wp.to_torch(cube_2.data.root_quat_w), - wp.to_torch(cube_3.data.root_quat_w), + cube_1.data.root_quat_w.torch, + cube_2.data.root_quat_w.torch, + cube_3.data.root_quat_w.torch, ), dim=1, ) @@ -103,15 +103,9 @@ def instance_randomize_cube_orientations_in_world_frame( cube_2_quat_w = [] cube_3_quat_w = [] for env_id in range(env.num_envs): - cube_1_quat_w.append( - wp.to_torch(cube_1.data.body_link_quat_w)[env_id, env.rigid_objects_in_focus[env_id][0], :4] - ) - cube_2_quat_w.append( - wp.to_torch(cube_2.data.body_link_quat_w)[env_id, env.rigid_objects_in_focus[env_id][1], :4] - ) - cube_3_quat_w.append( - wp.to_torch(cube_3.data.body_link_quat_w)[env_id, env.rigid_objects_in_focus[env_id][2], :4] - ) + cube_1_quat_w.append(cube_1.data.body_link_quat_w.torch[env_id, env.rigid_objects_in_focus[env_id][0], :4]) + cube_2_quat_w.append(cube_2.data.body_link_quat_w.torch[env_id, env.rigid_objects_in_focus[env_id][1], :4]) + cube_3_quat_w.append(cube_3.data.body_link_quat_w.torch[env_id, env.rigid_objects_in_focus[env_id][2], :4]) cube_1_quat_w = torch.stack(cube_1_quat_w) cube_2_quat_w = torch.stack(cube_2_quat_w) cube_3_quat_w = torch.stack(cube_3_quat_w) @@ -146,16 +140,16 @@ def object_obs( cube_3: RigidObject = env.scene[cube_3_cfg.name] ee_frame: FrameTransformer = env.scene[ee_frame_cfg.name] - cube_1_pos_w = wp.to_torch(cube_1.data.root_pos_w) - cube_1_quat_w = wp.to_torch(cube_1.data.root_quat_w) + cube_1_pos_w = cube_1.data.root_pos_w.torch + cube_1_quat_w = cube_1.data.root_quat_w.torch - cube_2_pos_w = wp.to_torch(cube_2.data.root_pos_w) - cube_2_quat_w = wp.to_torch(cube_2.data.root_quat_w) + cube_2_pos_w = cube_2.data.root_pos_w.torch + cube_2_quat_w = cube_2.data.root_quat_w.torch - cube_3_pos_w = wp.to_torch(cube_3.data.root_pos_w) - cube_3_quat_w = wp.to_torch(cube_3.data.root_quat_w) + cube_3_pos_w = cube_3.data.root_pos_w.torch + cube_3_quat_w = cube_3.data.root_quat_w.torch - ee_pos_w = wp.to_torch(ee_frame.data.target_pos_w)[:, 0, :] + ee_pos_w = ee_frame.data.target_pos_w.torch[:, 0, :] gripper_to_cube_1 = cube_1_pos_w - ee_pos_w gripper_to_cube_2 = cube_2_pos_w - ee_pos_w gripper_to_cube_3 = cube_3_pos_w - ee_pos_w @@ -220,18 +214,12 @@ def instance_randomize_object_obs( cube_2_quat_w = [] cube_3_quat_w = [] for env_id in range(env.num_envs): - cube_1_pos_w.append(wp.to_torch(cube_1.data.body_link_pos_w)[env_id, env.rigid_objects_in_focus[env_id][0], :3]) - cube_2_pos_w.append(wp.to_torch(cube_2.data.body_link_pos_w)[env_id, env.rigid_objects_in_focus[env_id][1], :3]) - cube_3_pos_w.append(wp.to_torch(cube_3.data.body_link_pos_w)[env_id, env.rigid_objects_in_focus[env_id][2], :3]) - cube_1_quat_w.append( - wp.to_torch(cube_1.data.body_link_quat_w)[env_id, env.rigid_objects_in_focus[env_id][0], :4] - ) - cube_2_quat_w.append( - wp.to_torch(cube_2.data.body_link_quat_w)[env_id, env.rigid_objects_in_focus[env_id][1], :4] - ) - cube_3_quat_w.append( - wp.to_torch(cube_3.data.body_link_quat_w)[env_id, env.rigid_objects_in_focus[env_id][2], :4] - ) + cube_1_pos_w.append(cube_1.data.body_link_pos_w.torch[env_id, env.rigid_objects_in_focus[env_id][0], :3]) + cube_2_pos_w.append(cube_2.data.body_link_pos_w.torch[env_id, env.rigid_objects_in_focus[env_id][1], :3]) + cube_3_pos_w.append(cube_3.data.body_link_pos_w.torch[env_id, env.rigid_objects_in_focus[env_id][2], :3]) + cube_1_quat_w.append(cube_1.data.body_link_quat_w.torch[env_id, env.rigid_objects_in_focus[env_id][0], :4]) + cube_2_quat_w.append(cube_2.data.body_link_quat_w.torch[env_id, env.rigid_objects_in_focus[env_id][1], :4]) + cube_3_quat_w.append(cube_3.data.body_link_quat_w.torch[env_id, env.rigid_objects_in_focus[env_id][2], :4]) cube_1_pos_w = torch.stack(cube_1_pos_w) cube_2_pos_w = torch.stack(cube_2_pos_w) cube_3_pos_w = torch.stack(cube_3_pos_w) @@ -239,7 +227,7 @@ def instance_randomize_object_obs( cube_2_quat_w = torch.stack(cube_2_quat_w) cube_3_quat_w = torch.stack(cube_3_quat_w) - ee_pos_w = wp.to_torch(ee_frame.data.target_pos_w)[:, 0, :] + ee_pos_w = ee_frame.data.target_pos_w.torch[:, 0, :] gripper_to_cube_1 = cube_1_pos_w - ee_pos_w gripper_to_cube_2 = cube_2_pos_w - ee_pos_w gripper_to_cube_3 = cube_3_pos_w - ee_pos_w @@ -269,14 +257,14 @@ def instance_randomize_object_obs( def ee_frame_pos(env: ManagerBasedRLEnv, ee_frame_cfg: SceneEntityCfg = SceneEntityCfg("ee_frame")) -> torch.Tensor: ee_frame: FrameTransformer = env.scene[ee_frame_cfg.name] - ee_frame_pos = wp.to_torch(ee_frame.data.target_pos_w)[:, 0, :] - env.scene.env_origins[:, 0:3] + ee_frame_pos = ee_frame.data.target_pos_w.torch[:, 0, :] - env.scene.env_origins[:, 0:3] return ee_frame_pos def ee_frame_quat(env: ManagerBasedRLEnv, ee_frame_cfg: SceneEntityCfg = SceneEntityCfg("ee_frame")) -> torch.Tensor: ee_frame: FrameTransformer = env.scene[ee_frame_cfg.name] - ee_frame_quat = wp.to_torch(ee_frame.data.target_quat_w)[:, 0, :] + ee_frame_quat = ee_frame.data.target_quat_w.torch[:, 0, :] return ee_frame_quat @@ -305,8 +293,8 @@ def gripper_pos( if hasattr(env.cfg, "gripper_joint_names"): gripper_joint_ids, _ = robot.find_joints(env.cfg.gripper_joint_names) assert len(gripper_joint_ids) == 2, "Observation gripper_pos only support parallel gripper for now" - finger_joint_1 = wp.to_torch(robot.data.joint_pos)[:, gripper_joint_ids[0]].clone().unsqueeze(1) - finger_joint_2 = -1 * wp.to_torch(robot.data.joint_pos)[:, gripper_joint_ids[1]].clone().unsqueeze(1) + finger_joint_1 = robot.data.joint_pos.torch[:, gripper_joint_ids[0]].clone().unsqueeze(1) + finger_joint_2 = -1 * robot.data.joint_pos.torch[:, gripper_joint_ids[1]].clone().unsqueeze(1) return torch.cat((finger_joint_1, finger_joint_2), dim=1) else: raise NotImplementedError("[Error] Cannot find gripper_joint_names in the environment config") @@ -325,8 +313,8 @@ def object_grasped( ee_frame: FrameTransformer = env.scene[ee_frame_cfg.name] object: RigidObject = env.scene[object_cfg.name] - object_pos = wp.to_torch(object.data.root_pos_w) - end_effector_pos = wp.to_torch(ee_frame.data.target_pos_w)[:, 0, :] + object_pos = object.data.root_pos_w.torch + end_effector_pos = ee_frame.data.target_pos_w.torch[:, 0, :] pose_diff = torch.linalg.vector_norm(object_pos - end_effector_pos, dim=1) if hasattr(env.scene, "surface_grippers") and len(env.scene.surface_grippers) > 0: @@ -343,7 +331,7 @@ def object_grasped( grasped = torch.logical_and( pose_diff < diff_threshold, torch.abs( - wp.to_torch(robot.data.joint_pos)[:, gripper_joint_ids[0]] + robot.data.joint_pos.torch[:, gripper_joint_ids[0]] - torch.tensor(env.cfg.gripper_open_val, dtype=torch.float32).to(env.device) ) > env.cfg.gripper_threshold, @@ -351,7 +339,7 @@ def object_grasped( grasped = torch.logical_and( grasped, torch.abs( - wp.to_torch(robot.data.joint_pos)[:, gripper_joint_ids[1]] + robot.data.joint_pos.torch[:, gripper_joint_ids[1]] - torch.tensor(env.cfg.gripper_open_val, dtype=torch.float32).to(env.device) ) > env.cfg.gripper_threshold, @@ -375,7 +363,7 @@ def object_stacked( upper_object: RigidObject = env.scene[upper_object_cfg.name] lower_object: RigidObject = env.scene[lower_object_cfg.name] - pos_diff = wp.to_torch(upper_object.data.root_pos_w) - wp.to_torch(lower_object.data.root_pos_w) + pos_diff = upper_object.data.root_pos_w.torch - lower_object.data.root_pos_w.torch height_dist = torch.linalg.vector_norm(pos_diff[:, 2:], dim=1) xy_dist = torch.linalg.vector_norm(pos_diff[:, :2], dim=1) @@ -393,7 +381,7 @@ def object_stacked( assert len(gripper_joint_ids) == 2, "Observations only support parallel gripper for now" stacked = torch.logical_and( torch.isclose( - wp.to_torch(robot.data.joint_pos)[:, gripper_joint_ids[0]], + robot.data.joint_pos.torch[:, gripper_joint_ids[0]], torch.tensor(env.cfg.gripper_open_val, dtype=torch.float32).to(env.device), atol=1e-4, rtol=1e-4, @@ -402,7 +390,7 @@ def object_stacked( ) stacked = torch.logical_and( torch.isclose( - wp.to_torch(robot.data.joint_pos)[:, gripper_joint_ids[1]], + robot.data.joint_pos.torch[:, gripper_joint_ids[1]], torch.tensor(env.cfg.gripper_open_val, dtype=torch.float32).to(env.device), atol=1e-4, rtol=1e-4, @@ -429,17 +417,17 @@ def cube_poses_in_base_frame( cube_2: RigidObject = env.scene[cube_2_cfg.name] cube_3: RigidObject = env.scene[cube_3_cfg.name] - pos_cube_1_world = wp.to_torch(cube_1.data.root_pos_w) - pos_cube_2_world = wp.to_torch(cube_2.data.root_pos_w) - pos_cube_3_world = wp.to_torch(cube_3.data.root_pos_w) + pos_cube_1_world = cube_1.data.root_pos_w.torch + pos_cube_2_world = cube_2.data.root_pos_w.torch + pos_cube_3_world = cube_3.data.root_pos_w.torch - quat_cube_1_world = wp.to_torch(cube_1.data.root_quat_w) - quat_cube_2_world = wp.to_torch(cube_2.data.root_quat_w) - quat_cube_3_world = wp.to_torch(cube_3.data.root_quat_w) + quat_cube_1_world = cube_1.data.root_quat_w.torch + quat_cube_2_world = cube_2.data.root_quat_w.torch + quat_cube_3_world = cube_3.data.root_quat_w.torch robot: Articulation = env.scene[robot_cfg.name] - root_pos_w = wp.to_torch(robot.data.root_pos_w) - root_quat_w = wp.to_torch(robot.data.root_quat_w) + root_pos_w = robot.data.root_pos_w.torch + root_quat_w = robot.data.root_quat_w.torch pos_cube_1_base, quat_cube_1_base = math_utils.subtract_frame_transforms( root_pos_w, root_quat_w, pos_cube_1_world, quat_cube_1_world @@ -488,17 +476,17 @@ def object_abs_obs_in_base_frame( ee_frame: FrameTransformer = env.scene[ee_frame_cfg.name] robot: Articulation = env.scene[robot_cfg.name] - root_pos_w = wp.to_torch(robot.data.root_pos_w) - root_quat_w = wp.to_torch(robot.data.root_quat_w) + root_pos_w = robot.data.root_pos_w.torch + root_quat_w = robot.data.root_quat_w.torch - cube_1_pos_w = wp.to_torch(cube_1.data.root_pos_w) - cube_1_quat_w = wp.to_torch(cube_1.data.root_quat_w) + cube_1_pos_w = cube_1.data.root_pos_w.torch + cube_1_quat_w = cube_1.data.root_quat_w.torch - cube_2_pos_w = wp.to_torch(cube_2.data.root_pos_w) - cube_2_quat_w = wp.to_torch(cube_2.data.root_quat_w) + cube_2_pos_w = cube_2.data.root_pos_w.torch + cube_2_quat_w = cube_2.data.root_quat_w.torch - cube_3_pos_w = wp.to_torch(cube_3.data.root_pos_w) - cube_3_quat_w = wp.to_torch(cube_3.data.root_quat_w) + cube_3_pos_w = cube_3.data.root_pos_w.torch + cube_3_quat_w = cube_3.data.root_quat_w.torch pos_cube_1_base, quat_cube_1_base = math_utils.subtract_frame_transforms( root_pos_w, root_quat_w, cube_1_pos_w, cube_1_quat_w @@ -510,8 +498,8 @@ def object_abs_obs_in_base_frame( root_pos_w, root_quat_w, cube_3_pos_w, cube_3_quat_w ) - ee_pos_w = wp.to_torch(ee_frame.data.target_pos_w)[:, 0, :] - ee_quat_w = wp.to_torch(ee_frame.data.target_quat_w)[:, 0, :] + ee_pos_w = ee_frame.data.target_pos_w.torch[:, 0, :] + ee_quat_w = ee_frame.data.target_quat_w.torch[:, 0, :] ee_pos_base, ee_quat_base = math_utils.subtract_frame_transforms(root_pos_w, root_quat_w, ee_pos_w, ee_quat_w) return torch.cat( @@ -539,12 +527,12 @@ def ee_frame_pose_in_base_frame( The end effector pose in the robot base frame. """ ee_frame: FrameTransformer = env.scene[ee_frame_cfg.name] - ee_frame_pos_w = wp.to_torch(ee_frame.data.target_pos_w)[:, 0, :] - ee_frame_quat_w = wp.to_torch(ee_frame.data.target_quat_w)[:, 0, :] + ee_frame_pos_w = ee_frame.data.target_pos_w.torch[:, 0, :] + ee_frame_quat_w = ee_frame.data.target_quat_w.torch[:, 0, :] robot: Articulation = env.scene[robot_cfg.name] - root_pos_w = wp.to_torch(robot.data.root_pos_w) - root_quat_w = wp.to_torch(robot.data.root_quat_w) + root_pos_w = robot.data.root_pos_w.torch + root_quat_w = robot.data.root_quat_w.torch ee_pos_in_base, ee_quat_in_base = math_utils.subtract_frame_transforms( root_pos_w, root_quat_w, ee_frame_pos_w, ee_frame_quat_w ) diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/mdp/terminations.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/mdp/terminations.py index d8781b0367a7..e709821972cf 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/mdp/terminations.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/mdp/terminations.py @@ -39,7 +39,7 @@ def cubes_stacked( cube_1: RigidObject = env.scene[cube_1_cfg.name] cube_2: RigidObject = env.scene[cube_2_cfg.name] - pos_diff_c12 = wp.to_torch(cube_1.data.root_pos_w) - wp.to_torch(cube_2.data.root_pos_w) + pos_diff_c12 = cube_1.data.root_pos_w.torch - cube_2.data.root_pos_w.torch # Compute cube position difference in x-y plane xy_dist_c12 = torch.linalg.norm(pos_diff_c12[:, :2], dim=1) @@ -54,7 +54,7 @@ def cubes_stacked( if cube_3_cfg is not None: cube_3: RigidObject = env.scene[cube_3_cfg.name] - pos_diff_c23 = wp.to_torch(cube_2.data.root_pos_w) - wp.to_torch(cube_3.data.root_pos_w) + pos_diff_c23 = cube_2.data.root_pos_w.torch - cube_3.data.root_pos_w.torch # Compute cube position difference in x-y plane xy_dist_c23 = torch.linalg.norm(pos_diff_c23[:, :2], dim=1) @@ -81,7 +81,7 @@ def cubes_stacked( stacked = torch.logical_and( torch.isclose( - wp.to_torch(robot.data.joint_pos)[:, gripper_joint_ids[0]], + robot.data.joint_pos.torch[:, gripper_joint_ids[0]], torch.tensor(env.cfg.gripper_open_val, dtype=torch.float32).to(env.device), atol=atol, rtol=rtol, @@ -90,7 +90,7 @@ def cubes_stacked( ) stacked = torch.logical_and( torch.isclose( - wp.to_torch(robot.data.joint_pos)[:, gripper_joint_ids[1]], + robot.data.joint_pos.torch[:, gripper_joint_ids[1]], torch.tensor(env.cfg.gripper_open_val, dtype=torch.float32).to(env.device), atol=atol, rtol=rtol, diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/navigation/mdp/pre_trained_policy_action.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/navigation/mdp/pre_trained_policy_action.py index 43de222a506d..4857d63711e1 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/navigation/mdp/pre_trained_policy_action.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/navigation/mdp/pre_trained_policy_action.py @@ -8,7 +8,6 @@ from typing import TYPE_CHECKING import torch -import warp as wp import isaaclab.utils.math as math_utils from isaaclab.managers import ActionTerm, ObservationManager @@ -138,12 +137,12 @@ def _debug_vis_callback(self, event): return # get marker location # -- base state - base_pos_w = wp.to_torch(self.robot.data.root_pos_w).clone() + base_pos_w = self.robot.data.root_pos_w.torch.clone() base_pos_w[:, 2] += 0.5 # -- resolve the scales and quaternions vel_des_arrow_scale, vel_des_arrow_quat = self._resolve_xy_velocity_to_arrow(self.raw_actions[:, :2]) vel_arrow_scale, vel_arrow_quat = self._resolve_xy_velocity_to_arrow( - wp.to_torch(self.robot.data.root_lin_vel_b)[:, :2] + self.robot.data.root_lin_vel_b.torch[:, :2] ) # display markers self.base_vel_goal_visualizer.visualize(base_pos_w, vel_des_arrow_quat, vel_des_arrow_scale) @@ -165,7 +164,7 @@ def _resolve_xy_velocity_to_arrow(self, xy_velocity: torch.Tensor) -> tuple[torc zeros = torch.zeros_like(heading_angle) arrow_quat = math_utils.quat_from_euler_xyz(zeros, zeros, heading_angle) # convert everything back from base to world frame - base_quat_w = wp.to_torch(self.robot.data.root_quat_w) + base_quat_w = self.robot.data.root_quat_w.torch arrow_quat = math_utils.quat_mul(base_quat_w, arrow_quat) return arrow_scale, arrow_quat diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/cartpole/cartpole_warp_env.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/cartpole/cartpole_warp_env.py index a9670c6b8dd4..ff3c87da847e 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/cartpole/cartpole_warp_env.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/cartpole/cartpole_warp_env.py @@ -226,8 +226,8 @@ def __init__(self, cfg: CartpoleWarpEnvCfg, render_mode: str | None = None, **kw # Simulation bindings # Note: these are direct memory views into the Newton simulation data, they should not be modified directly - self.joint_pos = self.cartpole.data.joint_pos - self.joint_vel = self.cartpole.data.joint_vel + self.joint_pos = self.cartpole.data.joint_pos.warp + self.joint_vel = self.cartpole.data.joint_vel.warp # Buffers self.observations = wp.zeros((self.num_envs), dtype=wp.vec4f, device=self.device) @@ -344,8 +344,8 @@ def _reset_idx(self, mask: wp.array | None = None) -> None: reset, dim=self.num_envs, inputs=[ - self.cartpole.data.default_joint_pos, - self.cartpole.data.default_joint_vel, + self.cartpole.data.default_joint_pos.warp, + self.cartpole.data.default_joint_vel.warp, self.joint_pos, self.joint_vel, self._cart_dof_idx[0], diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/inhand_manipulation/inhand_manipulation_warp_env.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/inhand_manipulation/inhand_manipulation_warp_env.py index 8b142cf8e2ee..d5bc6a2ad529 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/inhand_manipulation/inhand_manipulation_warp_env.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/inhand_manipulation/inhand_manipulation_warp_env.py @@ -573,8 +573,8 @@ def __init__(self, cfg: AllegroHandWarpEnvCfg, render_mode: str | None = None, * self.finger_bodies = wp.array(finger_bodies, dtype=wp.int32, device=self.device) # joint limits - self.hand_dof_lower_limits = self.hand.data.joint_pos_limits_lower - self.hand_dof_upper_limits = self.hand.data.joint_pos_limits_upper + self.hand_dof_lower_limits = self.hand.data.joint_pos_limits_lower.warp + self.hand_dof_upper_limits = self.hand.data.joint_pos_limits_upper.warp # unit vectors self.x_unit_vec = wp.vec3f(1.0, 0.0, 0.0) @@ -603,7 +603,7 @@ def __init__(self, cfg: AllegroHandWarpEnvCfg, render_mode: str | None = None, * self.goal_pos_w = wp.zeros(self.num_envs, dtype=wp.vec3f, device=self.device) # Initialize goal constants from Torch (avoid a one-off kernel launch). - default_root_pose = wp.to_torch(self.object.data.default_root_pose).to(self.device) + default_root_pose = self.object.data.default_root_pose.torch.to(self.device) in_hand_pos = default_root_pose[:, 0:3].clone() in_hand_pos[:, 2] -= 0.04 self.in_hand_pos.assign(wp.from_torch(in_hand_pos, dtype=wp.vec3f)) @@ -812,15 +812,15 @@ def _reset_idx(self, mask: wp.array | None = None): reset_object, dim=self.num_envs, inputs=[ - self.object.data.default_root_pose, + self.object.data.default_root_pose.warp, self.env_origins, self.cfg.reset_position_noise, self.x_unit_vec, self.y_unit_vec, mask, self.rng_state, - self.object.data.root_link_pose_w, - self.object.data.root_com_vel_w, + self.object.data.root_link_pose_w.warp, + self.object.data.root_com_vel_w.warp, ], device=self.device, ) @@ -830,8 +830,8 @@ def _reset_idx(self, mask: wp.array | None = None): reset_hand, dim=self.num_envs, inputs=[ - self.hand.data.default_joint_pos, - self.hand.data.default_joint_vel, + self.hand.data.default_joint_pos.warp, + self.hand.data.default_joint_vel.warp, self.hand_dof_lower_limits, self.hand_dof_upper_limits, self.cfg.reset_dof_pos_noise, @@ -839,8 +839,8 @@ def _reset_idx(self, mask: wp.array | None = None): mask, self.num_hand_dofs, self.rng_state, - self.hand.data.joint_pos, - self.hand.data.joint_vel, + self.hand.data.joint_pos.warp, + self.hand.data.joint_vel.warp, self.prev_targets, self.cur_targets, self.hand_dof_targets, @@ -900,13 +900,13 @@ def _compute_intermediate_values(self): compute_intermediate_values, dim=self.num_envs, inputs=[ - self.hand.data.body_pos_w, - self.hand.data.body_quat_w, - self.hand.data.body_vel_w, + self.hand.data.body_pos_w.warp, + self.hand.data.body_quat_w.warp, + self.hand.data.body_vel_w.warp, self.finger_bodies, self.env_origins, - self.object.data.root_link_pose_w, - self.object.data.root_com_vel_w, + self.object.data.root_link_pose_w.warp, + self.object.data.root_com_vel_w.warp, self.num_fingertips, self.fingertip_pos, self.fingertip_rot, @@ -950,8 +950,8 @@ def compute_full_observations(self): compute_full_observations, dim=self.num_envs, inputs=[ - self.hand.data.joint_pos, - self.hand.data.joint_vel, + self.hand.data.joint_pos.warp, + self.hand.data.joint_vel.warp, self.hand_dof_lower_limits, self.hand_dof_upper_limits, self.cfg.vel_obs_scale, diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/locomotion/locomotion_env_warp.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/locomotion/locomotion_env_warp.py index 6a726eb035ef..d8d712a655e9 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/locomotion/locomotion_env_warp.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/locomotion/locomotion_env_warp.py @@ -359,11 +359,11 @@ def __init__(self, cfg: DirectRLEnvCfg, render_mode: str | None = None, **kwargs # Simulation bindings # Note: these are direct memory views into the Newton simulation data, they should not be modified directly - self.joint_pos = self.robot.data.joint_pos - self.joint_vel = self.robot.data.joint_vel - self.root_pose_w = self.robot.data.root_pose_w - self.root_vel_w = self.robot.data.root_vel_w - self.soft_joint_pos_limits = self.robot.data.soft_joint_pos_limits + self.joint_pos = self.robot.data.joint_pos.warp + self.joint_vel = self.robot.data.joint_vel.warp + self.root_pose_w = self.robot.data.root_pose_w.warp + self.root_vel_w = self.robot.data.root_vel_w.warp + self.soft_joint_pos_limits = self.robot.data.soft_joint_pos_limits.warp # Buffers self.observations = wp.zeros( @@ -548,8 +548,8 @@ def _reset_idx(self, mask: wp.array | None = None): reset_root, dim=self.num_envs, inputs=[ - self.robot.data.default_root_pose, - self.robot.data.default_root_vel, + self.robot.data.default_root_pose.warp, + self.robot.data.default_root_vel.warp, self.env_origins, self.cfg.sim.dt, self.targets, @@ -564,8 +564,8 @@ def _reset_idx(self, mask: wp.array | None = None): reset_joints, dim=(self.num_envs, self.robot.num_joints), inputs=[ - self.robot.data.default_joint_pos, - self.robot.data.default_joint_vel, + self.robot.data.default_joint_pos.warp, + self.robot.data.default_joint_vel.warp, self.joint_pos, self.joint_vel, mask, diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/cartpole/mdp/rewards.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/cartpole/mdp/rewards.py index fbb426751fad..b83d2ea61706 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/cartpole/mdp/rewards.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/cartpole/mdp/rewards.py @@ -37,10 +37,10 @@ def _joint_pos_target_l2_kernel( def joint_pos_target_l2(env: ManagerBasedRLEnv, out, target: float, asset_cfg: SceneEntityCfg) -> None: """Penalize joint position deviation from a target value. Writes into ``out``.""" asset: Articulation = env.scene[asset_cfg.name] - assert asset.data.joint_pos.shape[1] == asset_cfg.joint_mask.shape[0] + assert asset.data.joint_pos.warp.shape[1] == asset_cfg.joint_mask.shape[0] wp.launch( kernel=_joint_pos_target_l2_kernel, dim=env.num_envs, - inputs=[asset.data.joint_pos, asset_cfg.joint_mask, out, target], + inputs=[asset.data.joint_pos.warp, asset_cfg.joint_mask, out, target], device=env.device, ) diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/humanoid/mdp/observations.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/humanoid/mdp/observations.py index 0210c407dc5f..23c3c779eab9 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/humanoid/mdp/observations.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/humanoid/mdp/observations.py @@ -55,7 +55,7 @@ def base_yaw_roll(env: ManagerBasedEnv, out, asset_cfg: SceneEntityCfg = SceneEn wp.launch( kernel=_base_yaw_roll_kernel, dim=env.num_envs, - inputs=[asset.data.root_quat_w, out], + inputs=[asset.data.root_quat_w.warp, out], device=env.device, ) @@ -83,7 +83,7 @@ def base_up_proj(env: ManagerBasedEnv, out, asset_cfg: SceneEntityCfg = SceneEnt wp.launch( kernel=_base_up_proj_kernel, dim=env.num_envs, - inputs=[asset.data.root_link_pose_w, asset.data.GRAVITY_VEC_W, out], + inputs=[asset.data.root_link_pose_w.warp, asset.data.GRAVITY_VEC_W.warp, out], device=env.device, ) @@ -128,7 +128,14 @@ def base_heading_proj( wp.launch( kernel=_base_heading_proj_kernel, dim=env.num_envs, - inputs=[asset.data.root_pos_w, asset.data.root_quat_w, target_pos[0], target_pos[1], target_pos[2], out], + inputs=[ + asset.data.root_pos_w.warp, + asset.data.root_quat_w.warp, + target_pos[0], + target_pos[1], + target_pos[2], + out, + ], device=env.device, ) @@ -174,6 +181,6 @@ def base_angle_to_target( wp.launch( kernel=_base_angle_to_target_kernel, dim=env.num_envs, - inputs=[asset.data.root_pos_w, asset.data.root_quat_w, target_pos[0], target_pos[1], out], + inputs=[asset.data.root_pos_w.warp, asset.data.root_quat_w.warp, target_pos[0], target_pos[1], out], device=env.device, ) diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/humanoid/mdp/rewards.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/humanoid/mdp/rewards.py index e45ca6cac4d1..f0e47dfbbf05 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/humanoid/mdp/rewards.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/humanoid/mdp/rewards.py @@ -60,7 +60,7 @@ def upright_posture_bonus( wp.launch( kernel=_upright_posture_bonus_kernel, dim=env.num_envs, - inputs=[asset.data.root_link_pose_w, asset.data.GRAVITY_VEC_W, threshold, out], + inputs=[asset.data.root_link_pose_w.warp, asset.data.GRAVITY_VEC_W.warp, threshold, out], device=env.device, ) @@ -102,7 +102,14 @@ def move_to_target_bonus( wp.launch( kernel=_move_to_target_bonus_kernel, dim=env.num_envs, - inputs=[asset.data.root_pos_w, asset.data.root_quat_w, target_pos[0], target_pos[1], threshold, out], + inputs=[ + asset.data.root_pos_w.warp, + asset.data.root_quat_w.warp, + target_pos[0], + target_pos[1], + threshold, + out, + ], device=env.device, ) @@ -172,7 +179,7 @@ def reset(self, env_mask: wp.array | None = None) -> None: dim=self.num_envs, inputs=[ env_mask, - asset.data.root_pos_w, + asset.data.root_pos_w.warp, self._target_pos[0], self._target_pos[1], self._target_pos[2], @@ -194,7 +201,7 @@ def __call__( wp.launch( kernel=_progress_reward_kernel, dim=env.num_envs, - inputs=[asset.data.root_pos_w, target_pos[0], target_pos[1], inv_dt, self.potentials, out], + inputs=[asset.data.root_pos_w.warp, target_pos[0], target_pos[1], inv_dt, self.potentials, out], device=env.device, ) @@ -256,8 +263,8 @@ def __call__( kernel=_joint_pos_limits_penalty_ratio_kernel, dim=env.num_envs, inputs=[ - asset.data.joint_pos, - asset.data.soft_joint_pos_limits, + asset.data.joint_pos.warp, + asset.data.soft_joint_pos_limits.warp, self._gear_ratio_scaled_wp, threshold, 1.0 / (1.0 - threshold), @@ -309,6 +316,6 @@ def __call__( wp.launch( kernel=_power_consumption_kernel, dim=env.num_envs, - inputs=[env.action_manager.action, asset.data.joint_vel, self._gear_ratio_scaled_wp, out], + inputs=[env.action_manager.action, asset.data.joint_vel.warp, self._gear_ratio_scaled_wp, out], device=env.device, ) diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/mdp/curriculums.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/mdp/curriculums.py index 1ed0a4c6f33c..01653af7d96b 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/mdp/curriculums.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/mdp/curriculums.py @@ -15,7 +15,6 @@ from typing import TYPE_CHECKING import torch -import warp as wp from isaaclab.assets import Articulation from isaaclab.managers import SceneEntityCfg @@ -32,7 +31,7 @@ def terrain_levels_vel( asset: Articulation = env.scene[asset_cfg.name] terrain: TerrainImporter = env.scene.terrain command = env.command_manager.get_command("base_velocity") - distance = torch.norm(wp.to_torch(asset.data.root_pos_w)[env_ids, :2] - env.scene.env_origins[env_ids, :2], dim=1) + distance = torch.norm(asset.data.root_pos_w.torch[env_ids, :2] - env.scene.env_origins[env_ids, :2], dim=1) move_up = distance > terrain.cfg.terrain_generator.size[0] / 2 move_down = distance < torch.norm(command[env_ids, :2], dim=1) * env.max_episode_length_s * 0.5 move_down *= ~move_up diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/mdp/rewards.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/mdp/rewards.py index 5bf6de7a414b..6daea9d45499 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/mdp/rewards.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/mdp/rewards.py @@ -63,7 +63,14 @@ def feet_air_time(env: ManagerBasedRLEnv, out, command_name: str, sensor_cfg: Sc wp.launch( kernel=_feet_air_time_kernel, dim=env.num_envs, - inputs=[contact_sensor.data.last_air_time, first_contact, sensor_cfg.body_ids_wp, fn._cmd_wp, threshold, out], + inputs=[ + contact_sensor.data.last_air_time.warp, + first_contact, + sensor_cfg.body_ids_wp, + fn._cmd_wp, + threshold, + out, + ], device=env.device, ) @@ -118,8 +125,8 @@ def feet_air_time_positive_biped(env, out, command_name: str, threshold: float, kernel=_feet_air_time_positive_biped_kernel, dim=env.num_envs, inputs=[ - contact_sensor.data.current_air_time, - contact_sensor.data.current_contact_time, + contact_sensor.data.current_air_time.warp, + contact_sensor.data.current_contact_time.warp, sensor_cfg.body_ids_wp, fn._cmd_wp, threshold, @@ -169,10 +176,10 @@ def feet_slide(env, out, sensor_cfg: SceneEntityCfg, asset_cfg: SceneEntityCfg = kernel=_feet_slide_kernel, dim=env.num_envs, inputs=[ - asset.data.body_lin_vel_w, - contact_sensor.data.net_forces_w_history, + asset.data.body_lin_vel_w.warp, + contact_sensor.data.net_forces_w_history.warp, sensor_cfg.body_ids_wp, - contact_sensor.data.net_forces_w_history.shape[1], + contact_sensor.data.net_forces_w_history.warp.shape[1], out, ], device=env.device, @@ -227,7 +234,7 @@ def track_lin_vel_xy_yaw_frame_exp( wp.launch( kernel=_track_lin_vel_xy_yaw_frame_exp_kernel, dim=env.num_envs, - inputs=[asset.data.root_quat_w, asset.data.root_lin_vel_w, fn._cmd_wp, 1.0 / (std * std), out], + inputs=[asset.data.root_quat_w.warp, asset.data.root_lin_vel_w.warp, fn._cmd_wp, 1.0 / (std * std), out], device=env.device, ) @@ -263,7 +270,7 @@ def track_ang_vel_z_world_exp( wp.launch( kernel=_track_ang_vel_z_world_exp_kernel, dim=env.num_envs, - inputs=[asset.data.root_ang_vel_w, fn._cmd_wp, 1.0 / (std * std), out], + inputs=[asset.data.root_ang_vel_w.warp, fn._cmd_wp, 1.0 / (std * std), out], device=env.device, ) @@ -307,6 +314,6 @@ def stand_still_joint_deviation_l1( wp.launch( kernel=_stand_still_joint_deviation_l1_kernel, dim=env.num_envs, - inputs=[asset.data.joint_pos, asset.data.default_joint_pos, fn._cmd_wp, command_threshold, out], + inputs=[asset.data.joint_pos.warp, asset.data.default_joint_pos.warp, fn._cmd_wp, command_threshold, out], device=env.device, ) diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/mdp/terminations.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/mdp/terminations.py index d7405098973f..4b3699e0e0bf 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/mdp/terminations.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/mdp/terminations.py @@ -61,6 +61,6 @@ def terrain_out_of_bounds( wp.launch( kernel=_terrain_out_of_bounds_kernel, dim=env.num_envs, - inputs=[asset.data.root_pos_w, fn._half_width, fn._half_height, distance_buffer, out], + inputs=[asset.data.root_pos_w.warp, fn._half_width, fn._half_height, distance_buffer, out], device=env.device, ) diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/mdp/rewards.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/mdp/rewards.py index 811163ec9737..e378bbfcee80 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/mdp/rewards.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/mdp/rewards.py @@ -60,9 +60,9 @@ def position_command_error(env: ManagerBasedRLEnv, out, command_name: str, asset kernel=_position_command_error_kernel, dim=env.num_envs, inputs=[ - asset.data.root_pos_w, - asset.data.root_quat_w, - asset.data.body_pos_w, + asset.data.root_pos_w.warp, + asset.data.root_quat_w.warp, + asset.data.body_pos_w.warp, fn._cmd_wp, asset_cfg.body_ids[0], out, @@ -111,9 +111,9 @@ def position_command_error_tanh( kernel=_position_command_error_tanh_kernel, dim=env.num_envs, inputs=[ - asset.data.root_pos_w, - asset.data.root_quat_w, - asset.data.body_pos_w, + asset.data.root_pos_w.warp, + asset.data.root_quat_w.warp, + asset.data.body_pos_w.warp, fn._cmd_wp, asset_cfg.body_ids[0], 1.0 / std, @@ -161,6 +161,6 @@ def orientation_command_error(env: ManagerBasedRLEnv, out, command_name: str, as wp.launch( kernel=_orientation_command_error_kernel, dim=env.num_envs, - inputs=[asset.data.root_quat_w, asset.data.body_quat_w, fn._cmd_wp, asset_cfg.body_ids[0], out], + inputs=[asset.data.root_quat_w.warp, asset.data.body_quat_w.warp, fn._cmd_wp, asset_cfg.body_ids[0], out], device=env.device, ) diff --git a/source/isaaclab_teleop/test/test_oxr_device.py b/source/isaaclab_teleop/test/test_oxr_device.py index 1663a56612d9..be2a604a6489 100644 --- a/source/isaaclab_teleop/test/test_oxr_device.py +++ b/source/isaaclab_teleop/test/test_oxr_device.py @@ -183,9 +183,9 @@ def test_xr_anchor(empty_env, mock_xrcore): assert xr_anchor_view.count == 1 position, orientation = xr_anchor_view.get_world_poses() - np.testing.assert_almost_equal(position.numpy(), [[1, 2, 3]]) + np.testing.assert_almost_equal(position.torch.numpy(), [[1, 2, 3]]) # FrameView returns quaternion in xyzw format, identity is [0, 0, 0, 1] - np.testing.assert_almost_equal(orientation.numpy(), [[0, 0, 0, 1]]) + np.testing.assert_almost_equal(orientation.torch.numpy(), [[0, 0, 0, 1]]) # Check that xr anchor mode and custom anchor are set correctly assert carb.settings.get_settings().get("/persistent/xr/anchorMode") == "custom anchor" @@ -206,8 +206,8 @@ def test_xr_anchor_default(empty_env, mock_xrcore): assert xr_anchor_view.count == 1 position, orientation = xr_anchor_view.get_world_poses() - np.testing.assert_almost_equal(position.numpy().tolist(), [[0, 0, 0]]) - np.testing.assert_almost_equal(orientation.numpy().tolist(), [[0, 0, 0, 1]]) + np.testing.assert_almost_equal(position.torch.numpy().tolist(), [[0, 0, 0]]) + np.testing.assert_almost_equal(orientation.torch.numpy().tolist(), [[0, 0, 0, 1]]) # Check that xr anchor mode and custom anchor are set correctly assert carb.settings.get_settings().get("/persistent/xr/anchorMode") == "custom anchor" @@ -229,8 +229,8 @@ def test_xr_anchor_multiple_devices(empty_env, mock_xrcore): assert xr_anchor_view.count == 1 position, orientation = xr_anchor_view.get_world_poses() - np.testing.assert_almost_equal(position.numpy().tolist(), [[0, 0, 0]]) - np.testing.assert_almost_equal(orientation.numpy().tolist(), [[0, 0, 0, 1]]) + np.testing.assert_almost_equal(position.torch.numpy().tolist(), [[0, 0, 0]]) + np.testing.assert_almost_equal(orientation.torch.numpy().tolist(), [[0, 0, 0, 1]]) # Check that xr anchor mode and custom anchor are set correctly assert carb.settings.get_settings().get("/persistent/xr/anchorMode") == "custom anchor"