Optimizes getters of data inside asset classes (#2118)

# Description

Previously the root positions was obtained by performing the following
steps:

1. Call PhysX getters to get transforms and velocities
2. Concatenating the two to make a single state tensor
3. Indexing this tensor to get the position

For applications where we only want the positions, the above is an
expensive operation as we call CUDA-MemCpy twice and then do
concatenation which adds overhead.

This MR simplifies the above process by having separate lazy buffers for
poses, velocities and states.

Fixes # (issue)

<!-- As a practice, it is recommended to open an issue to have
discussions on the proposed pull request.
This makes it easier for the community to keep track of what is being
developed or added, and if a given feature
is demanded by more than one party. -->

## Type of change

- New feature (non-breaking change which adds functionality)
- Breaking change (fix or feature that would cause existing
functionality to not work as expected)
- This change requires a documentation update

## Screenshots

Please attach before and after screenshots of the change if applicable.

<!--
Example:

| Before | After |
| ------ | ----- |
| _gif/png before_ | _gif/png after_ |

To upload images to a PR -- simply drag and drop an image while in edit
mode and it should upload the image directly. You can then paste that
source into the above before/after sections.
-->

## Checklist

- [x] I have run the [`pre-commit` checks](https://pre-commit.com/) with
`./isaaclab.sh --format`
- [x] I have made corresponding changes to the documentation
- [x] My changes generate no new warnings
- [ ] I have added tests that prove my fix is effective or that my
feature works
- [x] I have updated the changelog and the corresponding version in the
extension's `config/extension.toml` file
- [x] I have added my name to the `CONTRIBUTORS.md` or my name already
exists there

---------
Signed-off-by: Kelly Guo <kellyg@nvidia.com>
Co-authored-by: Kelly Guo <kellyg@nvidia.com>

scripts/tutorials/01_assets/run_rigid_object.py

@@ -125,7 +125,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): {cone_object.data.root_state_w[:, :3]}")
+            print(f"Root position (in world): {cone_object.data.root_pos_w}")
 
 
 def main():

scripts/tutorials/05_controllers/run_diff_ik.py

@@ -165,8 +165,8 @@ def run_simulator(sim: sim_utils.SimulationContext, scene: InteractiveScene):
         else:
             # obtain quantities from simulation
             jacobian = robot.root_physx_view.get_jacobians()[:, ee_jacobi_idx, :, robot_entity_cfg.joint_ids]
-            ee_pose_w = robot.data.body_state_w[:, robot_entity_cfg.body_ids[0], 0:7]
-            root_pose_w = robot.data.root_state_w[:, 0:7]
+            ee_pose_w = robot.data.body_pose_w[:, robot_entity_cfg.body_ids[0]]
+            root_pose_w = robot.data.root_pose_w
             joint_pos = robot.data.joint_pos[:, robot_entity_cfg.joint_ids]
             # compute frame in root frame
             ee_pos_b, ee_quat_b = subtract_frame_transforms(

source/isaaclab/config/extension.toml

 [package]
 
 # Note: Semantic Versioning is used: https://semver.org/
-version = "0.40.2"
+version = "0.40.3"
 
 # Description
 title = "Isaac Lab framework for Robot Learning"

source/isaaclab/docs/CHANGELOG.rst

 Changelog
 ---------
 
+0.40.3 (2025-03-20)
+~~~~~~~~~~~~~~~~~~~
+
+Changed
+^^^^^^^
+
+* Made separate data buffers for poses and velocities for the :class:`~isaaclab.assets.Articulation`,
+  :class:`~isaaclab.assets.RigidObject`, and :class:`~isaaclab.assets.RigidObjectCollection` classes.
+  Previously, the two data buffers were stored together in a single buffer requiring an additional
+  concatenation operation when accessing the data.
+* Cleaned up ordering of members inside the data classes for the assets to make them easier
+  to comprehend. This reduced the code duplication within the class and made the class
+  more readable.
+
+
 0.40.2 (2025-05-10)
 ~~~~~~~~~~~~~~~~~~~
 

source/isaaclab/isaaclab/assets/deformable_object/deformable_object_data.py

@@ -127,10 +127,7 @@ class DeformableObjectData:
         Shape is (num_instances, max_sim_vertices_per_body, 6).
         """
         if self._nodal_state_w.timestamp < self._sim_timestamp:
-            nodal_positions = self.nodal_pos_w
-            nodal_velocities = self.nodal_vel_w
-            # set the buffer data and timestamp
-            self._nodal_state_w.data = torch.cat((nodal_positions, nodal_velocities), dim=-1)
+            self._nodal_state_w.data = torch.cat((self.nodal_pos_w, self.nodal_vel_w), dim=-1)
             self._nodal_state_w.timestamp = self._sim_timestamp
         return self._nodal_state_w.data
 

source/isaaclab/isaaclab/envs/mdp/commands/pose_command.py

@@ -106,8 +106,8 @@ class UniformPoseCommand(CommandTerm):
         pos_error, rot_error = compute_pose_error(
             self.pose_command_w[:, :3],
             self.pose_command_w[:, 3:],
-            self.robot.data.body_state_w[:, self.body_idx, :3],
-            self.robot.data.body_state_w[:, self.body_idx, 3:7],
+            self.robot.data.body_pos_w[:, self.body_idx],
+            self.robot.data.body_quat_w[:, self.body_idx],
         )
         self.metrics["position_error"] = torch.norm(pos_error, dim=-1)
         self.metrics["orientation_error"] = torch.norm(rot_error, dim=-1)
@@ -156,5 +156,5 @@ class UniformPoseCommand(CommandTerm):
         # -- goal pose
         self.goal_pose_visualizer.visualize(self.pose_command_w[:, :3], self.pose_command_w[:, 3:])
         # -- current body pose
-        body_link_state_w = self.robot.data.body_state_w[:, self.body_idx]
-        self.current_pose_visualizer.visualize(body_link_state_w[:, :3], body_link_state_w[:, 3:7])
+        body_link_pose_w = self.robot.data.body_link_pose_w[:, self.body_idx]
+        self.current_pose_visualizer.visualize(body_link_pose_w[:, :3], body_link_pose_w[:, 3:7])

source/isaaclab/isaaclab/envs/mdp/observations.py

@@ -150,7 +150,7 @@ def body_projected_gravity_b(
 
     body_quat = asset.data.body_quat_w[:, asset_cfg.body_ids]
     gravity_dir = asset.data.GRAVITY_VEC_W.unsqueeze(1)
-    return math_utils.quat_rotate_inverse(body_quat, gravity_dir).view(env.num_envs, -1)
+    return math_utils.quat_apply_inverse(body_quat, gravity_dir).view(env.num_envs, -1)
 
 
 """

source/isaaclab/isaaclab/scene/interactive_scene.py

@@ -543,7 +543,7 @@ class InteractiveScene:
         state["articulation"] = dict()
         for asset_name, articulation in self._articulations.items():
             asset_state = dict()
-            asset_state["root_pose"] = articulation.data.root_state_w[:, :7].clone()
+            asset_state["root_pose"] = articulation.data.root_pose_w.clone()
             if is_relative:
                 asset_state["root_pose"][:, :3] -= self.env_origins
             asset_state["root_velocity"] = articulation.data.root_vel_w.clone()
@@ -563,7 +563,7 @@ class InteractiveScene:
         state["rigid_object"] = dict()
         for asset_name, rigid_object in self._rigid_objects.items():
             asset_state = dict()
-            asset_state["root_pose"] = rigid_object.data.root_state_w[:, :7].clone()
+            asset_state["root_pose"] = rigid_object.data.root_pose_w.clone()
             if is_relative:
                 asset_state["root_pose"][:, :3] -= self.env_origins
             asset_state["root_velocity"] = rigid_object.data.root_vel_w.clone()

source/isaaclab/test/assets/test_articulation.py

@@ -1407,7 +1407,7 @@ def test_body_root_state(sim, num_articulations, device, with_offset):
             torch.testing.assert_close(pos_gt, body_com_state_w[..., :3], atol=1e-3, rtol=1e-1)
 
             # orientation
-            com_quat_b = articulation.data.com_quat_b
+            com_quat_b = articulation.data.body_com_quat_b
             com_quat_w = math_utils.quat_mul(body_link_state_w[..., 3:7], com_quat_b)
             torch.testing.assert_close(com_quat_w, body_com_state_w[..., 3:7])
             torch.testing.assert_close(com_quat_w[:, 0, :], root_com_state_w[..., 3:7])

source/isaaclab/test/assets/test_rigid_object.py

@@ -812,7 +812,7 @@ def test_body_root_state_properties(num_cubes, device, with_offset):
             else:
                 # cubes are spinning around center of mass
                 # position will not match
-                # center of mass position will be constant (i.e. spining around com)
+                # center of mass position will be constant (i.e. spinning around com)
                 torch.testing.assert_close(env_pos + offset, root_com_state_w[..., :3])
                 torch.testing.assert_close(env_pos + offset, body_com_state_w[..., :3].squeeze(-2))
                 # link position will be moving but should stay constant away from center of mass
@@ -828,7 +828,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 = cube_object.data.com_quat_b
+                com_quat_b = cube_object.data.body_com_quat_b
                 com_quat_w = quat_mul(body_link_state_w[..., 3:7], com_quat_b)
                 torch.testing.assert_close(com_quat_w, body_com_state_w[..., 3:7])
                 torch.testing.assert_close(com_quat_w.squeeze(-2), root_com_state_w[..., 3:7])
@@ -838,7 +838,7 @@ def test_body_root_state_properties(num_cubes, device, with_offset):
                 torch.testing.assert_close(body_state_w[..., 3:7], body_link_state_w[..., 3:7])
 
                 # lin_vel will not match
-                # center of mass vel will be constant (i.e. spining around com)
+                # center of mass vel will be constant (i.e. spinning around com)
                 torch.testing.assert_close(torch.zeros_like(root_com_state_w[..., 7:10]), root_com_state_w[..., 7:10])
                 torch.testing.assert_close(torch.zeros_like(body_com_state_w[..., 7:10]), body_com_state_w[..., 7:10])
                 # link frame will be moving, and should be equal to input angular velocity cross offset
@@ -883,7 +883,7 @@ def test_write_root_state(num_cubes, device, with_offset, state_location):
         com[..., :3] = offset.to("cpu")
         cube_object.root_physx_view.set_coms(com, env_idx)
 
-        # check ceter of mass has been set
+        # check center of mass has been set
         torch.testing.assert_close(cube_object.root_physx_view.get_coms(), com)
 
         rand_state = torch.zeros_like(cube_object.data.root_state_w)

source/isaaclab/test/assets/test_rigid_object_collection.py

@@ -402,7 +402,7 @@ def test_object_state_properties(sim, num_envs, num_cubes, device, with_offset,
 
     for i in range(10):
         # spin the object around Z axis (com)
-        cube_object.write_object_com_velocity_to_sim(spin_twist.repeat(num_envs, num_cubes, 1))
+        cube_object.write_object_velocity_to_sim(spin_twist.repeat(num_envs, num_cubes, 1))
         sim.step()
         cube_object.update(sim.cfg.dt)
 
@@ -430,7 +430,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 = cube_object.data.com_quat_b
+            com_quat_b = cube_object.data.object_com_quat_b
             com_quat_w = quat_mul(object_link_state_w[..., 3:7], com_quat_b)
             torch.testing.assert_close(com_quat_w, object_com_state_w[..., 3:7])
 

source/isaaclab/test/controllers/test_differential_ik.py

@@ -164,8 +164,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 = robot.data.body_state_w[:, ee_frame_idx, 0:7]
-    root_pose_w = robot.data.root_state_w[:, 0:7]
+    ee_pose_w = robot.data.body_pose_w[:, ee_frame_idx]
+    root_pose_w = robot.data.root_pose_w
     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]
     )
@@ -212,8 +212,8 @@ def _run_ik_controller(
             # so we MUST skip the first step
             # obtain quantities from simulation
             jacobian = robot.root_physx_view.get_jacobians()[:, ee_jacobi_idx, :, arm_joint_ids]
-            ee_pose_w = robot.data.body_state_w[:, ee_frame_idx, 0:7]
-            root_pose_w = robot.data.root_state_w[:, 0:7]
+            ee_pose_w = robot.data.body_pose_w[:, ee_frame_idx]
+            root_pose_w = robot.data.root_pose_w
             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, :])

source/isaaclab/test/controllers/test_operational_space.py

@@ -1416,8 +1416,8 @@ def _update_states(
     jacobian_b[:, 3:, :] = torch.bmm(root_rot_matrix, jacobian_b[:, 3:, :])
 
     # Compute current pose of the end-effector
-    root_pose_w = robot.data.root_state_w[:, 0:7]
-    ee_pose_w = robot.data.body_state_w[:, ee_frame_idx, 0:7]
+    root_pose_w = robot.data.root_pose_w
+    ee_pose_w = robot.data.body_pose_w[:, 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]
     )

source/isaaclab/test/sensors/test_frame_transformer.py

@@ -186,7 +186,7 @@ def test_frame_transformer_feet_wrt_base(sim):
 
         # check absolute frame transforms in world frame
         # -- ground-truth
-        root_pose_w = scene.articulations["robot"].data.root_state_w[:, :7]
+        root_pose_w = scene.articulations["robot"].data.root_pose_w
         feet_pos_w_gt = scene.articulations["robot"].data.body_pos_w[:, feet_indices]
         feet_quat_w_gt = scene.articulations["robot"].data.body_quat_w[:, feet_indices]
         # -- frame transformer
@@ -362,9 +362,9 @@ def test_frame_transformer_robot_body_to_external_cube(sim):
 
         # check absolute frame transforms in world frame
         # -- ground-truth
-        root_pose_w = scene.articulations["robot"].data.root_state_w[:, :7]
-        cube_pos_w_gt = scene.rigid_objects["cube"].data.root_state_w[:, :3]
-        cube_quat_w_gt = scene.rigid_objects["cube"].data.root_state_w[:, 3:7]
+        root_pose_w = scene.articulations["robot"].data.root_pose_w
+        cube_pos_w_gt = scene.rigid_objects["cube"].data.root_pos_w
+        cube_quat_w_gt = scene.rigid_objects["cube"].data.root_quat_w
         # -- frame transformer
         source_pos_w_tf = scene.sensors["frame_transformer"].data.source_pos_w
         source_quat_w_tf = scene.sensors["frame_transformer"].data.source_quat_w
@@ -451,8 +451,8 @@ def test_frame_transformer_offset_frames(sim):
 
         # check absolute frame transforms in world frame
         # -- ground-truth
-        cube_pos_w_gt = scene["cube"].data.root_state_w[:, :3]
-        cube_quat_w_gt = scene["cube"].data.root_state_w[:, 3:7]
+        cube_pos_w_gt = scene["cube"].data.root_pos_w
+        cube_quat_w_gt = scene["cube"].data.root_quat_w
         # -- frame transformer
         source_pos_w_tf = scene.sensors["frame_transformer"].data.source_pos_w
         source_quat_w_tf = scene.sensors["frame_transformer"].data.source_quat_w
@@ -544,7 +544,7 @@ def test_frame_transformer_all_bodies(sim):
 
         # check absolute frame transforms in world frame
         # -- ground-truth
-        root_pose_w = scene.articulations["robot"].data.root_state_w[:, :7]
+        root_pose_w = scene.articulations["robot"].data.root_pose_w
         bodies_pos_w_gt = scene.articulations["robot"].data.body_pos_w
         bodies_quat_w_gt = scene.articulations["robot"].data.body_quat_w
 

source/isaaclab_tasks/isaaclab_tasks/direct/quadcopter/quadcopter_env.py

@@ -157,7 +157,7 @@ class QuadcopterEnv(DirectRLEnv):
 
     def _get_observations(self) -> dict:
         desired_pos_b, _ = subtract_frame_transforms(
-            self._robot.data.root_state_w[:, :3], self._robot.data.root_state_w[:, 3:7], self._desired_pos_w
+            self._robot.data.root_pos_w, self._robot.data.root_quat_w, self._desired_pos_w
         )
         obs = torch.cat(
             [

source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/lift/mdp/observations.py

@@ -30,7 +30,5 @@ def object_position_in_robot_root_frame(
     robot: RigidObject = env.scene[robot_cfg.name]
     object: RigidObject = env.scene[object_cfg.name]
     object_pos_w = object.data.root_pos_w[:, :3]
-    object_pos_b, _ = subtract_frame_transforms(
-        robot.data.root_state_w[:, :3], robot.data.root_state_w[:, 3:7], object_pos_w
-    )
+    object_pos_b, _ = subtract_frame_transforms(robot.data.root_pos_w, robot.data.root_quat_w, object_pos_w)
     return object_pos_b

source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/lift/mdp/rewards.py

@@ -65,8 +65,8 @@ 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(robot.data.root_state_w[:, :3], robot.data.root_state_w[:, 3:7], des_pos_b)
+    des_pos_w, _ = combine_frame_transforms(robot.data.root_pos_w, robot.data.root_quat_w, des_pos_b)
     # distance of the end-effector to the object: (num_envs,)
-    distance = torch.norm(des_pos_w - object.data.root_pos_w[:, :3], dim=1)
+    distance = torch.norm(des_pos_w - object.data.root_pos_w, dim=1)
     # rewarded if the object is lifted above the threshold
     return (object.data.root_pos_w[:, 2] > minimal_height) * (1 - torch.tanh(distance / std))

source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/lift/mdp/terminations.py

@@ -50,7 +50,7 @@ def object_reached_goal(
     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(robot.data.root_state_w[:, :3], robot.data.root_state_w[:, 3:7], des_pos_b)
+    des_pos_w, _ = combine_frame_transforms(robot.data.root_pos_w, robot.data.root_quat_w, des_pos_b)
     # distance of the end-effector to the object: (num_envs,)
     distance = torch.norm(des_pos_w - object.data.root_pos_w[:, :3], dim=1)
 

source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/reach/mdp/rewards.py

@@ -33,8 +33,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(asset.data.root_state_w[:, :3], asset.data.root_state_w[:, 3:7], des_pos_b)
-    curr_pos_w = asset.data.body_state_w[:, asset_cfg.body_ids[0], :3]  # type: ignore
+    des_pos_w, _ = combine_frame_transforms(asset.data.root_pos_w, asset.data.root_quat_w, des_pos_b)
+    curr_pos_w = asset.data.body_pos_w[:, asset_cfg.body_ids[0]]  # type: ignore
     return torch.norm(curr_pos_w - des_pos_w, dim=1)
 
 
@@ -51,8 +51,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(asset.data.root_state_w[:, :3], asset.data.root_state_w[:, 3:7], des_pos_b)
-    curr_pos_w = asset.data.body_state_w[:, asset_cfg.body_ids[0], :3]  # type: ignore
+    des_pos_w, _ = combine_frame_transforms(asset.data.root_pos_w, asset.data.root_quat_w, des_pos_b)
+    curr_pos_w = asset.data.body_pos_w[:, asset_cfg.body_ids[0]]  # type: ignore
     distance = torch.norm(curr_pos_w - des_pos_w, dim=1)
     return 1 - torch.tanh(distance / std)
 
@@ -69,6 +69,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(asset.data.root_state_w[:, 3:7], des_quat_b)
-    curr_quat_w = asset.data.body_state_w[:, asset_cfg.body_ids[0], 3:7]  # type: ignore
+    des_quat_w = quat_mul(asset.data.root_quat_w, des_quat_b)
+    curr_quat_w = asset.data.body_quat_w[:, asset_cfg.body_ids[0]]  # type: ignore
     return quat_error_magnitude(curr_quat_w, des_quat_w)