Refactors pose and velocities to link frame and COM frame APIs

source/extensions/omni.isaac.lab/omni/isaac/lab/assets/articulation/articulation.py

@@ -308,10 +308,13 @@ def write_root_pose_to_sim(self, root_pose: torch.Tensor, env_ids: Sequence[int]
         self.root_physx_view.set_root_transforms(root_poses_xyzw, indices=physx_env_ids)
 
     def write_root_velocity_to_sim(self, root_velocity: torch.Tensor, env_ids: Sequence[int] | None = None):
-        """Set the root velocity over selected environment indices into the simulation.
+        """Set the root center of mass velocity over selected environment indices into the simulation.
+
+        The velocity comprises linear velocity (x, y, z) and angular velocity (x, y, z) in that order.
+        NOTE: This sets the velocity of the root's center of mass rather than the roots frame.
 
         Args:
-            root_velocity: Root velocities in simulation frame. Shape is (len(env_ids), 6).
+            root_velocity: Root center of mass velocities in simulation world frame. Shape is (len(env_ids), 6). 
             env_ids: Environment indices. If None, then all indices are used.
         """
         # resolve all indices
@@ -326,6 +329,26 @@ def write_root_velocity_to_sim(self, root_velocity: torch.Tensor, env_ids: Seque
         # set into simulation
         self.root_physx_view.set_root_velocities(self._data.root_state_w[:, 7:], indices=physx_env_ids)
 
+    def write_root_link_velocity_to_sim(self, root_velocity: torch.Tensor, env_ids: Sequence[int] | None = None):
+        """Set the root link velocity over selected environment indices into the simulation.
+
+        The velocity comprises linear velocity (x, y, z) and angular velocity (x, y, z) in that order.
+        NOTE: This sets the velocity of the root's frame rather than the roots center of mass.
+
+        Args:
+            root_velocity: Root frame velocities in simulation world frame. Shape is (len(env_ids), 6). 
+            env_ids: Environment indices. If None, then all indices are used.
+        """
+        root_com_velocity = root_velocity.clone()
+        quat = self._data.root_state_w[env_ids, :7]
+        com_pos_b = self._data._com_pos_b[env_ids,:]
+        # transform given velocity to center of mass
+        root_com_velocity[:,:3] += torch.linalg.cross(
+            root_com_velocity[:, 10:13], math_utils.quat_rotate(quat, com_pos_b), dim=-1
+        )
+        # write center of mass velocity to sim
+        self.write_root_velocity_to_sim(root_velocity=root_com_velocity,env_ids=env_ids)
+
     def write_joint_state_to_sim(
         self,
         position: torch.Tensor,

source/extensions/omni.isaac.lab/omni/isaac/lab/assets/articulation/articulation_data.py

@@ -70,6 +70,9 @@ def __init__(self, root_physx_view: physx.ArticulationView, device: str):
         self._joint_acc = TimestampedBuffer()
         self._joint_vel = TimestampedBuffer()
 
+        # Link center of mass
+        self._com_pos_b, _ = self._root_physx_view.get_coms().to(self.device).split([3, 4], dim=-1)
+
     def update(self, dt: float):
         # update the simulation timestamp
         self._sim_timestamp += dt
@@ -256,8 +259,8 @@ def update(self, dt: float):
     def root_state_w(self):
         """Root state ``[pos, quat, lin_vel, ang_vel]`` in simulation world frame. Shape is (num_instances, 13).
 
-        The position and quaternion are of the articulation root's actor frame. Meanwhile, the linear and angular
-        velocities are of the articulation root's center of mass frame.
+        The position and quaternion are of the articulation root's actor frame relative to the world. Meanwhile,
+        the linear and angular velocities are of the articulation root's center of mass frame.
         """
         if self._root_state_w.timestamp < self._sim_timestamp:
             # read data from simulation
@@ -269,6 +272,35 @@ def root_state_w(self):
             self._root_state_w.timestamp = self._sim_timestamp
         return self._root_state_w.data
 
+    @property
+    def root_state_link_w(self):
+        """Root state ``[pos, quat, lin_vel, ang_vel]`` in simulation world frame. Shape is (num_instances, 13).
+
+        The position, quaternion, and linear/angular velocity are of the articulation root's actor frame relative to the
+        world.
+        """
+        state = self.root_state_w.clone()
+        quat = state[:, 3:7]
+        # adjust linear velocity to link
+        state[:, 7:10] += torch.linalg.cross(
+            state[:, 10:13], math_utils.quat_rotate(quat, -self._com_pos_b[:, 0, :]), dim=-1
+        )
+        return state
+
+    @property
+    def root_state_com_w(self):
+        """Root state ``[pos, quat, lin_vel, ang_vel]`` in simulation world frame. Shape is (num_instances, 13).
+
+        The position, quaternion, and linear/angular velocity are of the articulation root link's center of mass frame
+        relative to the world. Center of mass frame is assumed to be the same orientation as the link rather than the
+        orientation of the principle inertia.
+        """
+        state = self.root_state_w.clone()
+        quat = state[:, 3:7]
+        # adjust position to center of mass
+        state[:, :3] += math_utils.quat_rotate(quat, self._com_pos_b[:, 0, :])
+        return state
+
     @property
     def body_state_w(self):
         """State of all bodies `[pos, quat, lin_vel, ang_vel]` in simulation world frame.
@@ -287,15 +319,46 @@ def body_state_w(self):
             self._body_state_w.timestamp = self._sim_timestamp
         return self._body_state_w.data
 
+    @property
+    def body_state_link_w(self):
+        """State of all bodies `[pos, quat, lin_vel, ang_vel]` in simulation world frame.
+        Shape is (num_instances, num_bodies, 13).
+
+        The position, quaternion, and linear/angular velocity are of the body's link frame relative to the world.
+        """
+        state = self.body_state_w.clone()
+        quat = state[..., 3:7]
+        # adjust linear velocity to link
+        state[..., 7:10] += torch.linalg.cross(
+            state[..., 10:13], math_utils.quat_rotate(quat, -self._com_pos_b), dim=-1
+        )
+        return state
+
+    @property
+    def body_state_com_w(self):
+        """State of all bodies `[pos, quat, lin_vel, ang_vel]` in simulation world frame.
+        Shape is (num_instances, num_bodies, 13).
+
+        The position, quaternion, and linear/angular velocity are of the body's center of mass frame relative to the
+        world. Center of mass frame is assumed to be the same orientation as the link rather than the orientation of the
+        principle inertia.
+        """
+        state = self.body_state_w.clone()
+        quat = state[..., 3:7]
+        # adjust position to center of mass
+        state[..., :3] -= math_utils.quat_rotate(quat, self._com_pos_b)
+        return state
+
     @property
     def body_acc_w(self):
-        """Acceleration of all bodies. Shape is (num_instances, num_bodies, 6).
+        """Acceleration of all bodies (center of mass). Shape is (num_instances, num_bodies, 6).
 
-        This quantity is the acceleration of the articulation links' center of mass frame.
+        All values are relative to the world.
         """
         if self._body_acc_w.timestamp < self._sim_timestamp:
             # read data from simulation and set the buffer data and timestamp
             self._body_acc_w.data = self._root_physx_view.get_link_accelerations()
+
             self._body_acc_w.timestamp = self._sim_timestamp
         return self._body_acc_w.data
 
@@ -353,113 +416,114 @@ def joint_acc(self):
     def root_pos_w(self) -> torch.Tensor:
         """Root position in simulation world frame. Shape is (num_instances, 3).
 
-        This quantity is the position of the actor frame of the articulation root.
+        This quantity is the position of the actor frame of the articulation root relative to the world.
         """
         return self.root_state_w[:, :3]
 
     @property
     def root_quat_w(self) -> torch.Tensor:
         """Root orientation (w, x, y, z) in simulation world frame. Shape is (num_instances, 4).
 
-        This quantity is the orientation of the actor frame of the articulation root.
+        This quantity is the orientation of the actor frame of the articulation root relative to the world.
         """
         return self.root_state_w[:, 3:7]
 
     @property
     def root_vel_w(self) -> torch.Tensor:
         """Root velocity in simulation world frame. Shape is (num_instances, 6).
 
-        This quantity contains the linear and angular velocities of the articulation root's center of
-        mass frame.
+        This quantity contains the linear and angular velocities of the articulation root's center of mass frame
+        relative to the world.
         """
         return self.root_state_w[:, 7:13]
 
     @property
     def root_lin_vel_w(self) -> torch.Tensor:
         """Root linear velocity in simulation world frame. Shape is (num_instances, 3).
 
-        This quantity is the linear velocity of the articulation root's center of mass frame.
+        This quantity is the linear velocity of the articulation root's center of mass frame relative to the world.
         """
         return self.root_state_w[:, 7:10]
 
     @property
     def root_ang_vel_w(self) -> torch.Tensor:
         """Root angular velocity in simulation world frame. Shape is (num_instances, 3).
 
-        This quantity is the angular velocity of the articulation root's center of mass frame.
+        This quantity is the angular velocity of the articulation root's center of mass frame relative to the world.
         """
         return self.root_state_w[:, 10:13]
 
     @property
     def root_lin_vel_b(self) -> torch.Tensor:
         """Root linear velocity in base frame. Shape is (num_instances, 3).
 
-        This quantity is the linear velocity of the articulation root's center of mass frame with
-        respect to the articulation root's actor frame.
+        This quantity is the linear velocity of the articulation root's center of mass frame relative to the world
+        with respect to the articulation root's actor frame.
         """
         return math_utils.quat_rotate_inverse(self.root_quat_w, self.root_lin_vel_w)
 
     @property
     def root_ang_vel_b(self) -> torch.Tensor:
         """Root angular velocity in base world frame. Shape is (num_instances, 3).
 
-        This quantity is the angular velocity of the articulation root's center of mass frame with respect to the
-        articulation root's actor frame.
+        This quantity is the angular velocity of the articulation root's center of mass frame relative to the world with
+        respect to the articulation root's actor frame.
         """
         return math_utils.quat_rotate_inverse(self.root_quat_w, self.root_ang_vel_w)
 
     @property
     def body_pos_w(self) -> torch.Tensor:
         """Positions of all bodies in simulation world frame. Shape is (num_instances, num_bodies, 3).
 
-        This quantity is the position of the rigid bodies' actor frame.
+        This quantity is the position of the rigid bodies' actor frame relative to the world.
         """
         return self.body_state_w[..., :3]
 
     @property
     def body_quat_w(self) -> torch.Tensor:
         """Orientation (w, x, y, z) of all bodies in simulation world frame. Shape is (num_instances, num_bodies, 4).
 
-        This quantity is the orientation of the rigid bodies' actor frame.
+        This quantity is the orientation of the rigid bodies' actor frame relative to the world.
         """
         return self.body_state_w[..., 3:7]
 
     @property
     def body_vel_w(self) -> torch.Tensor:
         """Velocity of all bodies in simulation world frame. Shape is (num_instances, num_bodies, 6).
 
-        This quantity contains the linear and angular velocities of the rigid bodies' center of mass frame.
+        This quantity contains the linear and angular velocities of the rigid bodies' center of mass frame relative
+        to the world.
         """
         return self.body_state_w[..., 7:13]
 
     @property
     def body_lin_vel_w(self) -> torch.Tensor:
         """Linear velocity of all bodies in simulation world frame. Shape is (num_instances, num_bodies, 3).
 
-        This quantity is the linear velocity of the rigid bodies' center of mass frame.
+        This quantity is the linear velocity of the rigid bodies' center of mass frame relative to the world.
         """
         return self.body_state_w[..., 7:10]
 
     @property
     def body_ang_vel_w(self) -> torch.Tensor:
         """Angular velocity of all bodies in simulation world frame. Shape is (num_instances, num_bodies, 3).
 
-        This quantity is the angular velocity of the rigid bodies' center of mass frame.
+        This quantity is the angular velocity of the rigid bodies' center of mass frame relative to the world.
         """
         return self.body_state_w[..., 10:13]
 
     @property
     def body_lin_acc_w(self) -> torch.Tensor:
         """Linear acceleration of all bodies in simulation world frame. Shape is (num_instances, num_bodies, 3).
 
-        This quantity is the linear acceleration of the rigid bodies' center of mass frame.
+        This quantity is the linear acceleration of the rigid bodies' center of mass frame relative to the world.
         """
         return self.body_acc_w[..., 0:3]
 
     @property
     def body_ang_acc_w(self) -> torch.Tensor:
         """Angular acceleration of all bodies in simulation world frame. Shape is (num_instances, num_bodies, 3).
 
-        This quantity is the angular acceleration of the rigid bodies' center of mass frame.
+        This quantity is the angular acceleration of the rigid bodies' center of mass frame relative to the world.
         """
         return self.body_acc_w[..., 3:6]

source/extensions/omni.isaac.lab/omni/isaac/lab/assets/rigid_object/rigid_object_data.py

@@ -66,6 +66,10 @@ def __init__(self, root_physx_view: physx.RigidBodyView, device: str):
         self._root_state_w = TimestampedBuffer()
         self._body_acc_w = TimestampedBuffer()
 
+        # Link center of mass
+        # com_pos_b, _ = self._root_physx_view.get_coms().to(self.device).split([3, 4], dim=-1)
+        self._com_pos_b, _ = self._root_physx_view.get_coms().to(self.device).split([3, 4], dim=-1)
+
     def update(self, dt: float):
         """Updates the data for the rigid object.
 
@@ -117,6 +121,33 @@ def root_state_w(self):
             self._root_state_w.timestamp = self._sim_timestamp
         return self._root_state_w.data
 
+    @property
+    def root_state_link_w(self):
+        """Root state ``[pos, quat, lin_vel, ang_vel]`` in simulation world frame. Shape is (num_instances, 13).
+
+        The position, quaternion, and linear/angular velocity are of the rigid body root frame relative to the
+        world.
+        """
+        state = self.root_state_w.clone()
+        quat = state[:, 3:7]
+        # adjust linear velocity to link
+        state[:, 7:10] += torch.linalg.cross(state[:, 10:13], math_utils.quat_rotate(quat, -self._com_pos_b), dim=-1)
+        return state
+
+    @property
+    def root_state_com_w(self):
+        """Root state ``[pos, quat, lin_vel, ang_vel]`` in simulation world frame. Shape is (num_instances, 13).
+
+        The position, quaternion, and linear/angular velocity are of the rigid body's center of mass frame
+        relative to the world. Center of mass frame is assumed to be the same orientation as the link rather than the
+        orientation of the principle inertia.
+        """
+        state = self.root_state_w.clone()
+        quat = state[:, 3:7]
+        # adjust position to center of mass
+        state[:, :3] += math_utils.quat_rotate(quat, self._com_pos_b)
+        return state
+
     @property
     def body_state_w(self):
         """State of all bodies `[pos, quat, lin_vel, ang_vel]` in simulation world frame. Shape is (num_instances, 1, 13).
@@ -126,6 +157,26 @@ def body_state_w(self):
         """
         return self.root_state_w.view(-1, 1, 13)
 
+    @property
+    def body_state_link_w(self):
+        """State of all bodies `[pos, quat, lin_vel, ang_vel]` in simulation world frame.
+        Shape is (num_instances,1, 13).
+
+        The position, quaternion, and linear/angular velocity are of the body's link frame relative to the world.
+        """
+        return self.root_state_link_w.view(-1, 1, 13)
+
+    @property
+    def body_state_com_w(self):
+        """State of all bodies `[pos, quat, lin_vel, ang_vel]` in simulation world frame.
+        Shape is (num_instances, num_bodies, 13).
+
+        The position, quaternion, and linear/angular velocity are of the body's center of mass frame relative to the
+        world. Center of mass frame is assumed to be the same orientation as the link rather than the orientation of the
+        principle inertia.
+        """
+        return self.root_state_com_w.view(-1, 1, 13)
+
     @property
     def body_acc_w(self):
         """Acceleration of all bodies. Shape is (num_instances, 1, 6).
@@ -186,7 +237,7 @@ def root_vel_w(self) -> torch.Tensor:
     def root_lin_vel_w(self) -> torch.Tensor:
         """Root linear velocity in simulation world frame. Shape is (num_instances, 3).
 
-        This quantity is the linear velocity of the root rigid body's center of mass frame.
+        This quantity is the linear velocity of the root rigid body's center of mass frame relative to the world.
         """
         return self.root_state_w[:, 7:10]