diff --git a/docs/source/migration/migrating_from_isaacgymenvs.rst b/docs/source/migration/migrating_from_isaacgymenvs.rst index 2632321917..22bc33c9a5 100644 --- a/docs/source/migration/migrating_from_isaacgymenvs.rst +++ b/docs/source/migration/migrating_from_isaacgymenvs.rst @@ -814,9 +814,9 @@ The ``progress_buf`` variable has also been renamed to ``episode_length_buf``. | | | | | self.joint_pos[env_ids] = joint_pos | | | | -| | self.cartpole.write_root_pose_to_sim( | +| | self.cartpole.write_root_link_pose_to_sim( | | | default_root_state[:, :7], env_ids) | -| | self.cartpole.write_root_velocity_to_sim( | +| | self.cartpole.write_root_com_velocity_to_sim( | | | default_root_state[:, 7:], env_ids) | | | self.cartpole.write_joint_state_to_sim( | | | joint_pos, joint_vel, None, env_ids) | diff --git a/docs/source/migration/migrating_from_omniisaacgymenvs.rst b/docs/source/migration/migrating_from_omniisaacgymenvs.rst index 7945ccb0b1..629be58bd7 100644 --- a/docs/source/migration/migrating_from_omniisaacgymenvs.rst +++ b/docs/source/migration/migrating_from_omniisaacgymenvs.rst @@ -364,8 +364,8 @@ In Isaac Lab, ``root_pose`` and ``root_velocity`` have been combined into single .. code-block::python - self.cartpole.write_root_pose_to_sim(default_root_state[:, :7], env_ids) - self.cartpole.write_root_velocity_to_sim(default_root_state[:, 7:], env_ids) + self.cartpole.write_root_link_pose_to_sim(default_root_state[:, :7], env_ids) + self.cartpole.write_root_com_velocity_to_sim(default_root_state[:, 7:], env_ids) Creating a New Environment @@ -738,9 +738,9 @@ reset the ``episode_length_buf`` buffer. | 1.0 - 2.0 * torch.rand(num_resets, device=self._device)) | self.joint_pos[env_ids] = joint_pos | | | self.joint_vel[env_ids] = joint_vel | | # apply resets | | -| indices = env_ids.to(dtype=torch.int32) | self.cartpole.write_root_pose_to_sim( | +| indices = env_ids.to(dtype=torch.int32) | self.cartpole.write_root_link_pose_to_sim( | | self._cartpoles.set_joint_positions(dof_pos, indices=indices) | default_root_state[:, :7], env_ids) | -| self._cartpoles.set_joint_velocities(dof_vel, indices=indices) | self.cartpole.write_root_velocity_to_sim( | +| self._cartpoles.set_joint_velocities(dof_vel, indices=indices) | self.cartpole.write_root_com_velocity_to_sim( | | | default_root_state[:, 7:], env_ids) | | # bookkeeping | self.cartpole.write_joint_state_to_sim( | | self.reset_buf[env_ids] = 0 | joint_pos, joint_vel, None, env_ids) | diff --git a/docs/source/overview/sensors/ray_caster.rst b/docs/source/overview/sensors/ray_caster.rst index b2846c9602..8c45d00c44 100644 --- a/docs/source/overview/sensors/ray_caster.rst +++ b/docs/source/overview/sensors/ray_caster.rst @@ -42,6 +42,7 @@ Using a ray caster sensor requires a **pattern** and a parent xform to be attach update_period = 1/60, offset=RayCasterCfg.OffsetCfg(pos=(0,0,0.5)), mesh_prim_paths=["/World/Ground"], + attach_yaw_only=True, pattern_cfg=patterns.LidarPatternCfg( channels = 100, vertical_fov_range=[-90, 90], diff --git a/docs/source/tutorials/01_assets/run_articulation.rst b/docs/source/tutorials/01_assets/run_articulation.rst index dc8d14c669..dd456d6c6d 100644 --- a/docs/source/tutorials/01_assets/run_articulation.rst +++ b/docs/source/tutorials/01_assets/run_articulation.rst @@ -66,8 +66,8 @@ Similar to a rigid object, an articulation also has a root state. This state cor articulation tree. On top of the root state, an articulation also has joint states. These states correspond to the joint positions and velocities. -To reset the articulation, we first set the root state by calling the :meth:`Articulation.write_root_state_to_sim` -method. Similarly, we set the joint states by calling the :meth:`Articulation.write_joint_state_to_sim` method. +To reset the articulation, we first set the root state by calling the :meth:`Articulation.write_root_link_pose_to_sim` and :meth:`Articulation.write_root_com_velocity_to_sim` +methods. Similarly, we set the joint states by calling the :meth:`Articulation.write_joint_state_to_sim` method. Finally, we call the :meth:`Articulation.reset` method to reset any internal buffers and caches. .. literalinclude:: ../../../../source/standalone/tutorials/01_assets/run_articulation.py diff --git a/docs/source/tutorials/01_assets/run_deformable_object.rst b/docs/source/tutorials/01_assets/run_deformable_object.rst index 4041d7e7ed..8edb2fdf4a 100644 --- a/docs/source/tutorials/01_assets/run_deformable_object.rst +++ b/docs/source/tutorials/01_assets/run_deformable_object.rst @@ -149,7 +149,7 @@ the average position of all the nodes in the mesh. .. literalinclude:: ../../../../source/standalone/tutorials/01_assets/run_deformable_object.py :language: python :start-at: # update buffers - :end-at: print(f"Root position (in world): {cube_object.data.root_pos_w[:, :3]}") + :end-at: print(f"Root position (in world): {cube_object.data.root_link_pos_w[:, :3]}") The Code Execution diff --git a/docs/source/tutorials/01_assets/run_rigid_object.rst b/docs/source/tutorials/01_assets/run_rigid_object.rst index f9fe141467..e4249c45e0 100644 --- a/docs/source/tutorials/01_assets/run_rigid_object.rst +++ b/docs/source/tutorials/01_assets/run_rigid_object.rst @@ -96,7 +96,7 @@ desired state of the rigid object prim into the world frame before setting it. We use the :attr:`assets.RigidObject.data.default_root_state` attribute to get the default root state of the spawned rigid object prims. This default state can be configured from the :attr:`assets.RigidObjectCfg.init_state` attribute, which we left as identity in this tutorial. We then randomize the translation of the root state and -set the desired state of the rigid object prim using the :meth:`assets.RigidObject.write_root_state_to_sim` method. +set the desired state of the rigid object prim using the :meth:`assets.RigidObject.write_root_link_pose_to_sim` and :meth:`assets.RigidObject.write_root_com_velocity_to_sim` methods. As the name suggests, this method writes the root state of the rigid object prim into the simulation buffer. .. literalinclude:: ../../../../source/standalone/tutorials/01_assets/run_rigid_object.py diff --git a/docs/source/tutorials/05_controllers/run_diff_ik.rst b/docs/source/tutorials/05_controllers/run_diff_ik.rst index 99600b0591..ab451ab8ca 100644 --- a/docs/source/tutorials/05_controllers/run_diff_ik.rst +++ b/docs/source/tutorials/05_controllers/run_diff_ik.rst @@ -79,7 +79,7 @@ joint positions, current end-effector pose, and the Jacobian matrix. While the attribute :attr:`assets.ArticulationData.joint_pos` provides the joint positions, we only want the joint positions of the robot's arm, and not the gripper. Similarly, while -the attribute :attr:`assets.ArticulationData.body_state_w` provides the state of all the +the attribute :attr:`assets.ArticulationData.body_link_state_w` provides the state of all the robot's bodies, we only want the state of the robot's end-effector. Thus, we need to index into these arrays to obtain the desired quantities. diff --git a/source/extensions/omni.isaac.lab/config/extension.toml b/source/extensions/omni.isaac.lab/config/extension.toml index 4b44b9a4bd..05f5be440a 100644 --- a/source/extensions/omni.isaac.lab/config/extension.toml +++ b/source/extensions/omni.isaac.lab/config/extension.toml @@ -1,7 +1,7 @@ [package] # Note: Semantic Versioning is used: https://semver.org/ -version = "0.29.3" +version = "0.30.0" # Description title = "Isaac Lab framework for Robot Learning" diff --git a/source/extensions/omni.isaac.lab/docs/CHANGELOG.rst b/source/extensions/omni.isaac.lab/docs/CHANGELOG.rst index 0dfdb83c2e..fe823cb371 100644 --- a/source/extensions/omni.isaac.lab/docs/CHANGELOG.rst +++ b/source/extensions/omni.isaac.lab/docs/CHANGELOG.rst @@ -1,6 +1,15 @@ Changelog --------- +0.30.0 (2024-12-16) +~~~~~~~~~~~~~~~~~~~ + +Changed +^^^^^^^ + +* Previously, physx returns the rigid bodies and articulations velocities in the com of bodies rather than the link frame, while poses are in link frames. We now explicitly provide :attr:`body_link_state` and :attr:`body_com_state` APIs replacing the previous :attr:`body_state` API. Previous APIs are now marked as deprecated. Please update any code using the previous pose and velocity APIs to use the new ``*_link_*`` or ``*_com_*`` APIs in :attr:`omni.isaac_lab.assets.RigidBody`, :attr:`omni.isaac_lab.assets.RigidBodyCollection`, and :attr:`omni.isaac_lab.assets.Articulation`. + + 0.29.3 (2024-12-16) ~~~~~~~~~~~~~~~~~~~ diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/assets/articulation/articulation.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/assets/articulation/articulation.py index ae3634e302..602a9efaa1 100644 --- a/source/extensions/omni.isaac.lab/omni/isaac/lab/assets/articulation/articulation.py +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/assets/articulation/articulation.py @@ -98,6 +98,10 @@ def __init__(self, cfg: ArticulationCfg): """ super().__init__(cfg) + self._root_state_dep_warn = False + self._root_pose_dep_warn = False + self._root_vel_dep_warn = False + """ Properties """ @@ -280,15 +284,71 @@ def write_root_state_to_sim(self, root_state: torch.Tensor, env_ids: Sequence[in root_state: Root state in simulation frame. Shape is (len(env_ids), 13). env_ids: Environment indices. If None, then all indices are used. """ + + # deprecation warning + if not self._root_state_dep_warn: + omni.log.warn( + "DeprecationWarning: Articluation.write_root_state_to_sim will be removed in a future release. Please" + " use write_root_link_state_to_sim or write_root_com_state_to_sim instead." + ) + self._root_state_dep_warn = True + # set into simulation self.write_root_pose_to_sim(root_state[:, :7], env_ids=env_ids) self.write_root_velocity_to_sim(root_state[:, 7:], env_ids=env_ids) + def write_root_com_state_to_sim(self, root_state: torch.Tensor, env_ids: Sequence[int] | None = None): + """Set the root center of mass state over selected environment indices into the simulation. + + The root state comprises of the cartesian position, quaternion orientation in (w, x, y, z), and linear + and angular velocity. All the quantities are in the simulation frame. + + Args: + root_state: Root state in simulation frame. Shape is (len(env_ids), 13). + env_ids: Environment indices. If None, then all indices are used. + """ + # set into simulation + self.write_root_com_pose_to_sim(root_state[:, :7], env_ids=env_ids) + self.write_root_com_velocity_to_sim(root_state[:, 7:], env_ids=env_ids) + + def write_root_link_state_to_sim(self, root_state: torch.Tensor, env_ids: Sequence[int] | None = None): + """Set the root link state over selected environment indices into the simulation. + + The root state comprises of the cartesian position, quaternion orientation in (w, x, y, z), and linear + and angular velocity. All the quantities are in the simulation frame. + + Args: + root_state: Root state in simulation frame. Shape is (len(env_ids), 13). + env_ids: Environment indices. If None, then all indices are used. + """ + # set into simulation + self.write_root_link_pose_to_sim(root_state[:, :7], env_ids=env_ids) + self.write_root_link_velocity_to_sim(root_state[:, 7:], env_ids=env_ids) + def write_root_pose_to_sim(self, root_pose: torch.Tensor, env_ids: Sequence[int] | None = None): """Set the root pose over selected environment indices into the simulation. The root pose comprises of the cartesian position and quaternion orientation in (w, x, y, z). + Args: + root_pose: Root poses in simulation frame. Shape is (len(env_ids), 7). + env_ids: Environment indices. If None, then all indices are used. + """ + # deprecation warning + if not self._root_pose_dep_warn: + omni.log.warn( + "DeprecationWarning: Articluation.write_root_pos_to_sim will be removed in a future release. Please use" + " write_root_link_pose_to_sim or write_root_com_pose_to_sim instead." + ) + self._root_pose_dep_warn = True + + self.write_root_link_pose_to_sim(root_pose, env_ids) + + def write_root_link_pose_to_sim(self, root_pose: torch.Tensor, env_ids: Sequence[int] | None = None): + """Set the root link pose over selected environment indices into the simulation. + + The root pose comprises of the cartesian position and quaternion orientation in (w, x, y, z). + Args: root_pose: Root poses in simulation frame. Shape is (len(env_ids), 7). env_ids: Environment indices. If None, then all indices are used. @@ -300,22 +360,78 @@ def write_root_pose_to_sim(self, root_pose: torch.Tensor, env_ids: Sequence[int] physx_env_ids = self._ALL_INDICES # note: we need to do this here since tensors are not set into simulation until step. # set into internal buffers - self._data.root_state_w[env_ids, :7] = root_pose.clone() + self._data.root_link_state_w[env_ids, :7] = root_pose.clone() + self._data._ignore_dep_warn = True + self._data.root_state_w[env_ids, :7] = self._data.root_link_state_w[env_ids, :7] + self._data._ignore_dep_warn = False # convert root quaternion from wxyz to xyzw - root_poses_xyzw = self._data.root_state_w[:, :7].clone() + root_poses_xyzw = self._data.root_link_state_w[:, :7].clone() root_poses_xyzw[:, 3:] = math_utils.convert_quat(root_poses_xyzw[:, 3:], to="xyzw") # Need to invalidate the buffer to trigger the update with the new root pose. self._data._body_state_w.timestamp = -1.0 + self._data._body_link_state_w.timestamp = -1.0 + self._data._body_com_state_w.timestamp = -1.0 # set into simulation self.root_physx_view.set_root_transforms(root_poses_xyzw, indices=physx_env_ids) + def write_root_com_pose_to_sim(self, root_pose: torch.Tensor, env_ids: Sequence[int] | None = None): + """Set the root center of mass pose over selected environment indices into the simulation. + + The root pose comprises of the cartesian position and quaternion orientation in (w, x, y, z). + The orientation is the orientation of the principle axes of inertia. + + Args: + root_pose: Root center of mass poses in simulation frame. Shape is (len(env_ids), 7). + env_ids: Environment indices. If None, then all indices are used. + """ + # resolve all indices + if env_ids is None: + local_env_ids = slice(None) + + com_pos = self.data.com_pos_b[local_env_ids, 0, :] + com_quat = self.data.com_quat_b[local_env_ids, 0, :] + + root_link_pos, root_link_quat = math_utils.combine_frame_transforms( + root_pose[..., :3], + root_pose[..., 3:7], + math_utils.quat_rotate(math_utils.quat_inv(com_quat), -com_pos), + math_utils.quat_inv(com_quat), + ) + + root_link_pose = torch.cat((root_link_pos, root_link_quat), dim=-1) + self.write_root_link_pose_to_sim(root_pose=root_link_pose, env_ids=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 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. + """ + # deprecation warning + if not self._root_vel_dep_warn: + omni.log.warn( + "DeprecationWarning: Articluation.write_root_velocity_to_sim will be removed in a future release." + " Please use write_root_link_velocity_to_sim or write_root_com_velocity_to_sim instead." + ) + self._root_vel_dep_warn = True + + self.write_root_com_velocity_to_sim(root_velocity=root_velocity, env_ids=env_ids) + + def write_root_com_velocity_to_sim(self, root_velocity: torch.Tensor, env_ids: Sequence[int] | None = None): + """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 physx_env_ids = env_ids if env_ids is None: @@ -323,10 +439,37 @@ def write_root_velocity_to_sim(self, root_velocity: torch.Tensor, env_ids: Seque physx_env_ids = self._ALL_INDICES # note: we need to do this here since tensors are not set into simulation until step. # set into internal buffers - self._data.root_state_w[env_ids, 7:] = root_velocity.clone() + self._data.root_com_state_w[env_ids, 7:] = root_velocity.clone() + self._data._ignore_dep_warn = True + self._data.root_state_w[env_ids, 7:] = self._data.root_com_state_w[env_ids, 7:] + self._data._ignore_dep_warn = False self._data.body_acc_w[env_ids] = 0.0 # set into simulation - self.root_physx_view.set_root_velocities(self._data.root_state_w[:, 7:], indices=physx_env_ids) + self.root_physx_view.set_root_velocities(self._data.root_com_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. + """ + # resolve all indices + if env_ids is None: + local_env_ids = slice(None) + + root_com_velocity = root_velocity.clone() + quat = self.data.root_link_state_w[local_env_ids, 3:7] + com_pos_b = self.data.com_pos_b[local_env_ids, 0, :] + # transform given velocity to center of mass + root_com_velocity[:, :3] += torch.linalg.cross( + root_com_velocity[:, 3:], math_utils.quat_rotate(quat, com_pos_b), dim=-1 + ) + # write center of mass velocity to sim + self.write_root_com_velocity_to_sim(root_velocity=root_com_velocity, env_ids=env_ids) def write_joint_state_to_sim( self, @@ -360,6 +503,8 @@ def write_joint_state_to_sim( self._data.joint_acc[env_ids, joint_ids] = 0.0 # Need to invalidate the buffer to trigger the update with the new root pose. self._data._body_state_w.timestamp = -1.0 + self._data._body_link_state_w.timestamp = -1.0 + self._data._body_com_state_w.timestamp = -1.0 # set into simulation self.root_physx_view.set_dof_positions(self._data.joint_pos, indices=physx_env_ids) self.root_physx_view.set_dof_velocities(self._data.joint_vel, indices=physx_env_ids) diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/assets/articulation/articulation_data.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/assets/articulation/articulation_data.py index 633785f2c2..a53b88ad6d 100644 --- a/source/extensions/omni.isaac.lab/omni/isaac/lab/assets/articulation/articulation_data.py +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/assets/articulation/articulation_data.py @@ -6,6 +6,7 @@ import torch import weakref +import omni.log import omni.physics.tensors.impl.api as physx import omni.isaac.lab.utils.math as math_utils @@ -64,12 +65,21 @@ def __init__(self, root_physx_view: physx.ArticulationView, device: str): # Initialize the lazy buffers. self._root_state_w = TimestampedBuffer() + self._root_link_state_w = TimestampedBuffer() + self._root_com_state_w = TimestampedBuffer() self._body_state_w = TimestampedBuffer() + self._body_link_state_w = TimestampedBuffer() + self._body_com_state_w = TimestampedBuffer() self._body_acc_w = TimestampedBuffer() self._joint_pos = TimestampedBuffer() self._joint_acc = TimestampedBuffer() self._joint_vel = TimestampedBuffer() + # deprecation warning check + self._root_state_dep_warn = False + self._body_state_dep_warn = False + self._ignore_dep_warn = False + def update(self, dt: float): # update the simulation timestamp self._sim_timestamp += dt @@ -211,12 +221,6 @@ def update(self, dt: float): joint_damping: torch.Tensor = None """Joint damping provided to simulation. Shape is (num_instances, num_joints).""" - joint_armature: torch.Tensor = None - """Joint armature provided to simulation. Shape is (num_instances, num_joints).""" - - joint_friction: torch.Tensor = None - """Joint friction provided to simulation. Shape is (num_instances, num_joints).""" - joint_limits: torch.Tensor = None """Joint limits provided to simulation. Shape is (num_instances, num_joints, 2).""" @@ -266,9 +270,17 @@ 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 not self._root_state_dep_warn and not self._ignore_dep_warn: + omni.log.warn( + "DeprecationWarning: root_state_w and it's derived properties will be deprecated in a future release." + " Please use root_link_state_w or root_com_state_w." + ) + self._root_state_dep_warn = True + if self._root_state_w.timestamp < self._sim_timestamp: # read data from simulation pose = self._root_physx_view.get_root_transforms().clone() @@ -279,6 +291,53 @@ def root_state_w(self): self._root_state_w.timestamp = self._sim_timestamp return self._root_state_w.data + @property + def root_link_state_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. + """ + if self._root_link_state_w.timestamp < self._sim_timestamp: + # read data from simulation + pose = self._root_physx_view.get_root_transforms().clone() + pose[:, 3:7] = math_utils.convert_quat(pose[:, 3:7], to="wxyz") + velocity = self._root_physx_view.get_root_velocities().clone() + + # adjust linear velocity to link from center of mass + velocity[:, :3] += torch.linalg.cross( + velocity[:, 3:], math_utils.quat_rotate(pose[:, 3:7], -self.com_pos_b[:, 0, :]), dim=-1 + ) + # set the buffer data and timestamp + self._root_link_state_w.data = torch.cat((pose, velocity), dim=-1) + self._root_link_state_w.timestamp = self._sim_timestamp + + return self._root_link_state_w.data + + @property + def root_com_state_w(self): + """Root center of mass 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. + """ + if self._root_com_state_w.timestamp < self._sim_timestamp: + # read data from simulation (pose is of link) + pose = self._root_physx_view.get_root_transforms().clone() + pose[:, 3:7] = math_utils.convert_quat(pose[:, 3:7], to="wxyz") + velocity = self._root_physx_view.get_root_velocities() + + # adjust pose to center of mass + pos, quat = math_utils.combine_frame_transforms( + pose[:, :3], pose[:, 3:7], self.com_pos_b[:, 0, :], self.com_quat_b[:, 0, :] + ) + pose = torch.cat((pos, quat), dim=-1) + # set the buffer data and timestamp + self._root_com_state_w.data = torch.cat((pose, velocity), dim=-1) + self._root_com_state_w.timestamp = self._sim_timestamp + return self._root_com_state_w.data + @property def body_state_w(self): """State of all bodies `[pos, quat, lin_vel, ang_vel]` in simulation world frame. @@ -287,6 +346,14 @@ def body_state_w(self): The position and quaternion are of all the articulation links's actor frame. Meanwhile, the linear and angular velocities are of the articulation links's center of mass frame. """ + + if not self._body_state_dep_warn and not self._ignore_dep_warn: + omni.log.warn( + "DeprecationWarning: body_state_w and it's derived properties will be deprecated in a future release." + " Please use body_link_state_w or body_com_state_w." + ) + self._body_state_dep_warn = True + if self._body_state_w.timestamp < self._sim_timestamp: self._physics_sim_view.update_articulations_kinematic() # read data from simulation @@ -298,22 +365,73 @@ def body_state_w(self): self._body_state_w.timestamp = self._sim_timestamp return self._body_state_w.data + @property + def body_link_state_w(self): + """State of all bodies' link frame`[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. + """ + if self._body_link_state_w.timestamp < self._sim_timestamp: + self._physics_sim_view.update_articulations_kinematic() + # read data from simulation + pose = self._root_physx_view.get_link_transforms().clone() + pose[..., 3:7] = math_utils.convert_quat(pose[..., 3:7], to="wxyz") + velocity = self._root_physx_view.get_link_velocities() + + # adjust linear velocity to link from center of mass + velocity[..., :3] += torch.linalg.cross( + velocity[..., 3:], math_utils.quat_rotate(pose[..., 3:7], -self.com_pos_b), dim=-1 + ) + # set the buffer data and timestamp + self._body_link_state_w.data = torch.cat((pose, velocity), dim=-1) + self._body_link_state_w.timestamp = self._sim_timestamp + + return self._body_link_state_w.data + + @property + def body_com_state_w(self): + """State of all bodies center of mass `[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. + """ + if self._body_com_state_w.timestamp < self._sim_timestamp: + self._physics_sim_view.update_articulations_kinematic() + # read data from simulation (pose is of link) + pose = self._root_physx_view.get_link_transforms().clone() + pose[..., 3:7] = math_utils.convert_quat(pose[..., 3:7], to="wxyz") + velocity = self._root_physx_view.get_link_velocities() + + # adjust pose to center of mass + pos, quat = math_utils.combine_frame_transforms( + pose[..., :3], pose[..., 3:7], self.com_pos_b, self.com_quat_b + ) + pose = torch.cat((pos, quat), dim=-1) + # set the buffer data and timestamp + self._body_com_state_w.data = torch.cat((pose, velocity), dim=-1) + self._body_com_state_w.timestamp = self._sim_timestamp + return self._body_com_state_w.data + @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 @property def projected_gravity_b(self): """Projection of the gravity direction on base frame. Shape is (num_instances, 3).""" - return math_utils.quat_rotate_inverse(self.root_quat_w, self.GRAVITY_VEC_W) + return math_utils.quat_rotate_inverse(self.root_link_quat_w, self.GRAVITY_VEC_W) @property def heading_w(self): @@ -323,7 +441,7 @@ def heading_w(self): This quantity is computed by assuming that the forward-direction of the base frame is along x-direction, i.e. :math:`(1, 0, 0)`. """ - forward_w = math_utils.quat_apply(self.root_quat_w, self.FORWARD_VEC_B) + forward_w = math_utils.quat_apply(self.root_link_quat_w, self.FORWARD_VEC_B) return torch.atan2(forward_w[:, 1], forward_w[:, 0]) @property @@ -364,7 +482,7 @@ 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] @@ -372,7 +490,7 @@ def root_pos_w(self) -> torch.Tensor: 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] @@ -380,8 +498,8 @@ def root_quat_w(self) -> torch.Tensor: 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] @@ -389,7 +507,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 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] @@ -397,7 +515,7 @@ def root_lin_vel_w(self) -> torch.Tensor: 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] @@ -405,8 +523,8 @@ def root_ang_vel_w(self) -> torch.Tensor: 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) @@ -414,16 +532,163 @@ def root_lin_vel_b(self) -> torch.Tensor: 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) + # + # Derived Root Link Frame Properties + # + + @property + def root_link_pos_w(self) -> torch.Tensor: + """Root link position in simulation world frame. Shape is (num_instances, 3). + + This quantity is the position of the actor frame of the root rigid body relative to the world. + """ + if self._body_com_state_w.timestamp < self._sim_timestamp: + # read data from simulation (pose is of link) + pose = self._root_physx_view.get_root_transforms() + return pose[:, :3] + return self.root_link_state_w[:, :3] + + @property + def root_link_quat_w(self) -> torch.Tensor: + """Root link 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 root rigid body. + """ + if self._body_com_state_w.timestamp < self._sim_timestamp: + # read data from simulation (pose is of link) + pose = self._root_physx_view.get_root_transforms().clone() + pose[:, 3:7] = math_utils.convert_quat(pose[:, 3:7], to="wxyz") + return pose[:, 3:7] + return self.root_link_state_w[:, 3:7] + + @property + def root_link_vel_w(self) -> torch.Tensor: + """Root link velocity in simulation world frame. Shape is (num_instances, 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_state_w[:, 7:13] + + @property + def root_link_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 actor frame relative to the world. + """ + return self.root_link_state_w[:, 7:10] + + @property + def root_link_ang_vel_w(self) -> torch.Tensor: + """Root link angular velocity in simulation world frame. Shape is (num_instances, 3). + + This quantity is the angular velocity of the actor frame of the root rigid body relative to the world. + """ + return self.root_link_state_w[:, 10:13] + + @property + def root_link_lin_vel_b(self) -> torch.Tensor: + """Root link linear velocity in base frame. Shape is (num_instances, 3). + + This quantity is the linear velocity of the actor frame of the root rigid body frame with respect to the + rigid body's actor frame. + """ + return math_utils.quat_rotate_inverse(self.root_link_quat_w, self.root_link_lin_vel_w) + + @property + def root_link_ang_vel_b(self) -> torch.Tensor: + """Root link angular velocity in base world frame. Shape is (num_instances, 3). + + This quantity is the angular velocity of the actor frame of the root rigid body frame with respect to the + rigid body's actor frame. + """ + return math_utils.quat_rotate_inverse(self.root_link_quat_w, self.root_link_ang_vel_w) + + # + # Root Center of Mass state properties + # + + @property + def root_com_pos_w(self) -> torch.Tensor: + """Root center of mass position in simulation world frame. Shape is (num_instances, 3). + + This quantity is the position of the actor frame of the root rigid body relative to the world. + """ + return self.root_com_state_w[:, :3] + + @property + def root_com_quat_w(self) -> torch.Tensor: + """Root center of mass 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 root rigid body relative to the world. + """ + return self.root_com_state_w[:, 3:7] + + @property + def root_com_vel_w(self) -> torch.Tensor: + """Root center of mass velocity in simulation world frame. Shape is (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. + """ + if self._root_com_state_w.timestamp < self._sim_timestamp: + # read data from simulation (pose is of link) + velocity = self._root_physx_view.get_root_velocities() + return velocity + return self.root_com_state_w[:, 7:13] + + @property + def root_com_lin_vel_w(self) -> torch.Tensor: + """Root center of mass 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 relative to the world. + """ + if self._root_com_state_w.timestamp < self._sim_timestamp: + # read data from simulation (pose is of link) + velocity = self._root_physx_view.get_root_velocities() + return velocity[:, 0:3] + return self.root_com_state_w[:, 7:10] + + @property + def root_com_ang_vel_w(self) -> torch.Tensor: + """Root center of mass angular velocity in simulation world frame. Shape is (num_instances, 3). + + This quantity is the angular velocity of the root rigid body's center of mass frame relative to the world. + """ + if self._root_com_state_w.timestamp < self._sim_timestamp: + self._physics_sim_view.update_articulations_kinematic() + # read data from simulation (pose is of link) + velocity = self._root_physx_view.get_root_velocities() + return velocity[:, 3:6] + return self.root_com_state_w[:, 10:13] + + @property + def root_com_lin_vel_b(self) -> torch.Tensor: + """Root center of mass linear velocity in base frame. Shape is (num_instances, 3). + + This quantity is the linear velocity of the root rigid body's center of mass frame with respect to the + rigid body's actor frame. + """ + return math_utils.quat_rotate_inverse(self.root_link_quat_w, self.root_com_lin_vel_w) + + @property + def root_com_ang_vel_b(self) -> torch.Tensor: + """Root center of mass angular velocity in base world frame. Shape is (num_instances, 3). + + This quantity is the angular velocity of the root rigid body's center of mass frame with respect to the + rigid body's actor frame. + """ + return math_utils.quat_rotate_inverse(self.root_link_quat_w, self.root_com_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] @@ -431,7 +696,7 @@ def body_pos_w(self) -> torch.Tensor: 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] @@ -439,7 +704,8 @@ def body_quat_w(self) -> torch.Tensor: 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] @@ -447,7 +713,7 @@ def body_vel_w(self) -> torch.Tensor: 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] @@ -455,7 +721,7 @@ def body_lin_vel_w(self) -> torch.Tensor: 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] @@ -463,7 +729,7 @@ def body_ang_vel_w(self) -> torch.Tensor: 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] @@ -471,6 +737,138 @@ def body_lin_acc_w(self) -> torch.Tensor: 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] + + # + # Link body properties + # + + @property + def body_link_pos_w(self) -> torch.Tensor: + """Positions of all bodies in simulation world frame. Shape is (num_instances, 1, 3). + + This quantity is the position of the rigid bodies' actor frame relative to the world. + """ + if self._body_link_state_w.timestamp < self._sim_timestamp: + self._physics_sim_view.update_articulations_kinematic() + # read data from simulation + pose = self._root_physx_view.get_link_transforms() + return pose[..., :3] + return self._body_link_state_w.data[..., :3] + + @property + def body_link_quat_w(self) -> torch.Tensor: + """Orientation (w, x, y, z) of all bodies in simulation world frame. Shape is (num_instances, 1, 4). + + This quantity is the orientation of the rigid bodies' actor frame relative to the world. + """ + if self._body_link_state_w.timestamp < self._sim_timestamp: + self._physics_sim_view.update_articulations_kinematic() + # read data from simulation + pose = self._root_physx_view.get_link_transforms().clone() + pose[..., 3:7] = math_utils.convert_quat(pose[..., 3:7], to="wxyz") + return pose[..., 3:7] + return self.body_link_state_w[..., 3:7] + + @property + def body_link_vel_w(self) -> torch.Tensor: + """Velocity of all bodies in simulation world frame. Shape is (num_instances, 1, 6). + + This quantity contains the linear and angular velocities of the rigid bodies' center of mass frame + relative to the world. + """ + return self.body_link_state_w[..., 7:13] + + @property + def body_link_lin_vel_w(self) -> torch.Tensor: + """Linear velocity of all bodies in simulation world frame. Shape is (num_instances, 1, 3). + + This quantity is the linear velocity of the rigid bodies' center of mass frame relative to the world. + """ + return self.body_link_state_w[..., 7:10] + + @property + def body_link_ang_vel_w(self) -> torch.Tensor: + """Angular velocity of all bodies in simulation world frame. Shape is (num_instances, 1, 3). + + This quantity is the angular velocity of the rigid bodies' center of mass frame relative to the world. + """ + return self.body_link_state_w[..., 10:13] + + # + # Center of mass body properties + # + + @property + def body_com_pos_w(self) -> torch.Tensor: + """Positions of all bodies in simulation world frame. Shape is (num_instances, 1, 3). + + This quantity is the position of the rigid bodies' actor frame. + """ + return self.body_com_state_w[..., :3] + + @property + def body_com_quat_w(self) -> torch.Tensor: + """Orientation (w, x, y, z) of the prinicple axies of inertia of all bodies in simulation world frame. + + Shape is (num_instances, 1, 4). This quantity is the orientation of the rigid bodies' actor frame. + """ + return self.body_com_state_w[..., 3:7] + + @property + def body_com_vel_w(self) -> torch.Tensor: + """Velocity of all bodies in simulation world frame. Shape is (num_instances, 1, 6). + + This quantity contains the linear and angular velocities of the rigid bodies' center of mass frame. + """ + if self._body_com_state_w.timestamp < self._sim_timestamp: + self._physics_sim_view.update_articulations_kinematic() + # read data from simulation (velocity is of com) + velocity = self._root_physx_view.get_link_velocities() + return velocity + return self.body_com_state_w[..., 7:13] + + @property + def body_com_lin_vel_w(self) -> torch.Tensor: + """Linear velocity of all bodies in simulation world frame. Shape is (num_instances, 1, 3). + + This quantity is the linear velocity of the rigid bodies' center of mass frame. + """ + if self._body_com_state_w.timestamp < self._sim_timestamp: + self._physics_sim_view.update_articulations_kinematic() + # read data from simulation (velocity is of com) + velocity = self._root_physx_view.get_link_velocities() + return velocity[..., 0:3] + return self.body_com_state_w[..., 7:10] + + @property + def body_com_ang_vel_w(self) -> torch.Tensor: + """Angular velocity of all bodies in simulation world frame. Shape is (num_instances, 1, 3). + + This quantity is the angular velocity of the rigid bodies' center of mass frame. + """ + if self._body_com_state_w.timestamp < self._sim_timestamp: + self._physics_sim_view.update_articulations_kinematic() + # read data from simulation (velocity is of com) + velocity = self._root_physx_view.get_link_velocities() + return velocity[..., 3:6] + return self.body_com_state_w[..., 10:13] + + @property + def com_pos_b(self) -> torch.Tensor: + """Center of mass of all of the bodies in simulation world frame. Shape is (num_instances, num_bodies, 3). + + This quantity is the center of mass location relative to its body frame. + """ + return self._root_physx_view.get_coms().to(self.device)[..., :3] + + @property + def com_quat_b(self) -> torch.Tensor: + """Orientation (w,x,y,z) of the prinicple axies of inertia of all of the bodies in simulation world frame. Shape is (num_instances, num_bodies, 4). + + This quantity is the orientation of the principles axes of inertia relative to its body frame. + """ + quat = self._root_physx_view.get_coms().to(self.device)[..., 3:7] + return math_utils.convert_quat(quat, to="wxyz") diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/assets/rigid_object/rigid_object.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/assets/rigid_object/rigid_object.py index 4908db4dc5..426a06fbf2 100644 --- a/source/extensions/omni.isaac.lab/omni/isaac/lab/assets/rigid_object/rigid_object.py +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/assets/rigid_object/rigid_object.py @@ -55,6 +55,9 @@ def __init__(self, cfg: RigidObjectCfg): cfg: A configuration instance. """ super().__init__(cfg) + self._root_state_dep_warn = False + self._root_pose_dep_warn = False + self._root_vel_dep_warn = False """ Properties @@ -156,15 +159,70 @@ def write_root_state_to_sim(self, root_state: torch.Tensor, env_ids: Sequence[in root_state: Root state in simulation frame. Shape is (len(env_ids), 13). env_ids: Environment indices. If None, then all indices are used. """ + # deprecation warning + if not self._root_state_dep_warn: + omni.log.warn( + "DeprecationWarning: RigidObject.write_root_state_to_sim will be removed in a future release. Please" + " use write_root_link_state_to_sim or write_root_com_state_to_sim instead." + ) + self._root_state_dep_warn = True + # set into simulation self.write_root_pose_to_sim(root_state[:, :7], env_ids=env_ids) self.write_root_velocity_to_sim(root_state[:, 7:], env_ids=env_ids) + def write_root_com_state_to_sim(self, root_state: torch.Tensor, env_ids: Sequence[int] | None = None): + """Set the root center of mass state over selected environment indices into the simulation. + + The root state comprises of the cartesian position, quaternion orientation in (w, x, y, z), and linear + and angular velocity. All the quantities are in the simulation frame. + + Args: + root_state: Root state in simulation frame. Shape is (len(env_ids), 13). + env_ids: Environment indices. If None, then all indices are used. + """ + # set into simulation + self.write_root_com_pose_to_sim(root_state[:, :7], env_ids=env_ids) + self.write_root_com_velocity_to_sim(root_state[:, 7:], env_ids=env_ids) + + def write_root_link_state_to_sim(self, root_state: torch.Tensor, env_ids: Sequence[int] | None = None): + """Set the root link state over selected environment indices into the simulation. + + The root state comprises of the cartesian position, quaternion orientation in (w, x, y, z), and linear + and angular velocity. All the quantities are in the simulation frame. + + Args: + root_state: Root state in simulation frame. Shape is (len(env_ids), 13). + env_ids: Environment indices. If None, then all indices are used. + """ + # set into simulation + self.write_root_link_pose_to_sim(root_state[:, :7], env_ids=env_ids) + self.write_root_link_velocity_to_sim(root_state[:, 7:], env_ids=env_ids) + def write_root_pose_to_sim(self, root_pose: torch.Tensor, env_ids: Sequence[int] | None = None): """Set the root pose over selected environment indices into the simulation. The root pose comprises of the cartesian position and quaternion orientation in (w, x, y, z). + Args: + root_pose: Root poses in simulation frame. Shape is (len(env_ids), 7). + env_ids: Environment indices. If None, then all indices are used. + """ + # deprecation warning + if not self._root_pose_dep_warn: + omni.log.warn( + "DeprecationWarning: RigidObject.write_root_pose_to_sim will be removed in a future release. Please use" + " write_root_link_pose_to_sim or write_root_com_pose_to_sim instead." + ) + self._root_pose_dep_warn = True + + self.write_root_link_pose_to_sim(root_pose, env_ids) + + def write_root_link_pose_to_sim(self, root_pose: torch.Tensor, env_ids: Sequence[int] | None = None): + """Set the root link pose over selected environment indices into the simulation. + + The root pose comprises of the cartesian position and quaternion orientation in (w, x, y, z). + Args: root_pose: Root poses in simulation frame. Shape is (len(env_ids), 7). env_ids: Environment indices. If None, then all indices are used. @@ -176,20 +234,75 @@ def write_root_pose_to_sim(self, root_pose: torch.Tensor, env_ids: Sequence[int] physx_env_ids = self._ALL_INDICES # note: we need to do this here since tensors are not set into simulation until step. # set into internal buffers - self._data.root_state_w[env_ids, :7] = root_pose.clone() + self._data.root_link_state_w[env_ids, :7] = root_pose.clone() + self._data._ignore_dep_warn = True + self._data.root_state_w[env_ids, :7] = self._data.root_link_state_w[env_ids, :7] + self._data._ignore_dep_warn = False # convert root quaternion from wxyz to xyzw - root_poses_xyzw = self._data.root_state_w[:, :7].clone() + root_poses_xyzw = self._data.root_link_state_w[:, :7].clone() root_poses_xyzw[:, 3:] = math_utils.convert_quat(root_poses_xyzw[:, 3:], to="xyzw") # set into simulation self.root_physx_view.set_transforms(root_poses_xyzw, indices=physx_env_ids) + def write_root_com_pose_to_sim(self, root_pose: torch.Tensor, env_ids: Sequence[int] | None = None): + """Set the root center of mass pose over selected environment indices into the simulation. + + The root pose comprises of the cartesian position and quaternion orientation in (w, x, y, z). + The orientation is the orientation of the principle axes of inertia. + + Args: + root_pose: Root center of mass poses in simulation frame. Shape is (len(env_ids), 7). + env_ids: Environment indices. If None, then all indices are used. + """ + + # resolve all indices + if env_ids is None: + local_env_ids = slice(None) + + com_pos = self.data.com_pos_b[local_env_ids, 0, :] + com_quat = self.data.com_quat_b[local_env_ids, 0, :] + + root_link_pos, root_link_quat = math_utils.combine_frame_transforms( + root_pose[..., :3], + root_pose[..., 3:7], + math_utils.quat_rotate(math_utils.quat_inv(com_quat), -com_pos), + math_utils.quat_inv(com_quat), + ) + + root_link_pose = torch.cat((root_link_pos, root_link_quat), dim=-1) + self.write_root_link_pose_to_sim(root_pose=root_link_pose, env_ids=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 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. + """ + # deprecation warning + if not self._root_vel_dep_warn: + omni.log.warn( + "DeprecationWarning: RigidObject.write_root_velocity_to_sim will be removed in a future release. Please" + " use write_root_link_velocity_to_sim or write_root_com_velocity_to_sim instead." + ) + self._root_vel_dep_warn = True + + self.write_root_com_velocity_to_sim(root_velocity=root_velocity, env_ids=env_ids) + + def write_root_com_velocity_to_sim(self, root_velocity: torch.Tensor, env_ids: Sequence[int] | None = None): + """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 physx_env_ids = env_ids if env_ids is None: @@ -197,10 +310,37 @@ def write_root_velocity_to_sim(self, root_velocity: torch.Tensor, env_ids: Seque physx_env_ids = self._ALL_INDICES # note: we need to do this here since tensors are not set into simulation until step. # set into internal buffers - self._data.root_state_w[env_ids, 7:] = root_velocity.clone() + self._data.root_com_state_w[env_ids, 7:] = root_velocity.clone() + self._data._ignore_dep_warn = True + self._data.root_state_w[env_ids, 7:] = self._data.root_com_state_w[env_ids, 7:] + self._data._ignore_dep_warn = False self._data.body_acc_w[env_ids] = 0.0 # set into simulation - self.root_physx_view.set_velocities(self._data.root_state_w[:, 7:], indices=physx_env_ids) + self.root_physx_view.set_velocities(self._data.root_com_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. + """ + # resolve all indices + if env_ids is None: + local_env_ids = slice(None) + + root_com_velocity = root_velocity.clone() + quat = self.data.root_link_state_w[local_env_ids, 3:7] + com_pos_b = self.data.com_pos_b[local_env_ids, 0, :] + # transform given velocity to center of mass + root_com_velocity[:, :3] += torch.linalg.cross( + root_com_velocity[:, 3:], math_utils.quat_rotate(quat, com_pos_b), dim=-1 + ) + # write center of mass velocity to sim + self.write_root_com_velocity_to_sim(root_velocity=root_com_velocity, env_ids=env_ids) """ Operations - Setters. diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/assets/rigid_object/rigid_object_data.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/assets/rigid_object/rigid_object_data.py index 63d7be943d..dfb343064b 100644 --- a/source/extensions/omni.isaac.lab/omni/isaac/lab/assets/rigid_object/rigid_object_data.py +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/assets/rigid_object/rigid_object_data.py @@ -6,6 +6,7 @@ import torch import weakref +import omni.log import omni.physics.tensors.impl.api as physx import omni.isaac.lab.utils.math as math_utils @@ -64,8 +65,14 @@ def __init__(self, root_physx_view: physx.RigidBodyView, device: str): # Initialize the lazy buffers. self._root_state_w = TimestampedBuffer() + self._root_link_state_w = TimestampedBuffer() + self._root_com_state_w = TimestampedBuffer() self._body_acc_w = TimestampedBuffer() + # deprecation warning check + self._root_state_dep_warn = False + self._ignore_dep_warn = False + def update(self, dt: float): """Updates the data for the rigid object. @@ -114,6 +121,13 @@ def root_state_w(self): The position and orientation are of the rigid body's actor frame. Meanwhile, the linear and angular velocities are of the rigid body's center of mass frame. """ + if not self._root_state_dep_warn and not self._ignore_dep_warn: + omni.log.warn( + "DeprecationWarning: root_state_w and it's derived properties will be deprecated in a future release." + " Please use root_link_state_w or root_com_state_w." + ) + self._root_state_dep_warn = True + if self._root_state_w.timestamp < self._sim_timestamp: # read data from simulation pose = self._root_physx_view.get_transforms().clone() @@ -124,6 +138,52 @@ def root_state_w(self): self._root_state_w.timestamp = self._sim_timestamp return self._root_state_w.data + @property + def root_link_state_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. + """ + if self._root_link_state_w.timestamp < self._sim_timestamp: + # read data from simulation + pose = self._root_physx_view.get_transforms().clone() + pose[:, 3:7] = math_utils.convert_quat(pose[:, 3:7], to="wxyz") + velocity = self._root_physx_view.get_velocities().clone() + + # adjust linear velocity to link from center of mass + velocity[:, :3] += torch.linalg.cross( + velocity[:, 3:], math_utils.quat_rotate(pose[:, 3:7], -self.com_pos_b[:, 0, :]), dim=-1 + ) + # set the buffer data and timestamp + self._root_link_state_w.data = torch.cat((pose, velocity), dim=-1) + self._root_link_state_w.timestamp = self._sim_timestamp + + return self._root_link_state_w.data + + @property + def root_com_state_w(self): + """Root center of mass 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 the orientation principle axes of inertia. + """ + if self._root_com_state_w.timestamp < self._sim_timestamp: + # read data from simulation (pose is of link) + pose = self._root_physx_view.get_transforms().clone() + pose[:, 3:7] = math_utils.convert_quat(pose[:, 3:7], to="wxyz") + velocity = self._root_physx_view.get_velocities() + + # adjust pose to center of mass + pos, quat = math_utils.combine_frame_transforms( + pose[:, :3], pose[:, 3:7], self.com_pos_b[:, 0, :], self.com_quat_b[:, 0, :] + ) + pose = torch.cat((pos, quat), dim=-1) + # set the buffer data and timestamp + self._root_com_state_w.data = torch.cat((pos, quat, velocity), dim=-1) + self._root_com_state_w.timestamp = self._sim_timestamp + return self._root_com_state_w.data + @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). @@ -131,8 +191,34 @@ def body_state_w(self): The position and orientation are of the rigid bodies' actor frame. Meanwhile, the linear and angular velocities are of the rigid bodies' center of mass frame. """ + + omni.log.warn( + "DeprecationWarning: body_state_w and it's derived properties will be deprecated in a future release." + " Please use body_link_state_w or body_com_state_w." + ) + return self.root_state_w.view(-1, 1, 13) + @property + def body_link_state_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_link_state_w.view(-1, 1, 13) + + @property + def body_com_state_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_com_state_w.view(-1, 1, 13) + @property def body_acc_w(self): """Acceleration of all bodies. Shape is (num_instances, 1, 6). @@ -148,7 +234,7 @@ def body_acc_w(self): @property def projected_gravity_b(self): """Projection of the gravity direction on base frame. Shape is (num_instances, 3).""" - return math_utils.quat_rotate_inverse(self.root_quat_w, self.GRAVITY_VEC_W) + return math_utils.quat_rotate_inverse(self.root_link_quat_w, self.GRAVITY_VEC_W) @property def heading_w(self): @@ -158,7 +244,7 @@ def heading_w(self): This quantity is computed by assuming that the forward-direction of the base frame is along x-direction, i.e. :math:`(1, 0, 0)`. """ - forward_w = math_utils.quat_apply(self.root_quat_w, self.FORWARD_VEC_B) + forward_w = math_utils.quat_apply(self.root_link_quat_w, self.FORWARD_VEC_B) return torch.atan2(forward_w[:, 1], forward_w[:, 0]) ## @@ -193,7 +279,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] @@ -212,7 +298,7 @@ def root_lin_vel_b(self) -> torch.Tensor: This quantity is the linear velocity of the root rigid body's center of mass frame with respect to the rigid body's actor frame. """ - return math_utils.quat_rotate_inverse(self.root_quat_w, self.root_lin_vel_w) + return math_utils.quat_rotate_inverse(self.root_link_quat_w, self.root_lin_vel_w) @property def root_ang_vel_b(self) -> torch.Tensor: @@ -221,7 +307,145 @@ def root_ang_vel_b(self) -> torch.Tensor: This quantity is the angular velocity of the root rigid body's center of mass frame with respect to the rigid body's actor frame. """ - return math_utils.quat_rotate_inverse(self.root_quat_w, self.root_ang_vel_w) + return math_utils.quat_rotate_inverse(self.root_link_quat_w, self.root_ang_vel_w) + + @property + def root_link_pos_w(self) -> torch.Tensor: + """Root link position in simulation world frame. Shape is (num_instances, 3). + + This quantity is the position of the actor frame of the root rigid body relative to the world. + """ + if self._root_link_state_w.timestamp < self._sim_timestamp: + # read data from simulation + pose = self._root_physx_view.get_transforms() + return pose[:, :3] + return self.root_link_state_w[:, :3] + + @property + def root_link_quat_w(self) -> torch.Tensor: + """Root link 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 root rigid body. + """ + if self._root_link_state_w.timestamp < self._sim_timestamp: + # read data from simulation + pose = self._root_physx_view.get_transforms().clone() + pose[:, 3:7] = math_utils.convert_quat(pose[:, 3:7], to="wxyz") + return pose[:, 3:7] + return self.root_link_state_w[:, 3:7] + + @property + def root_link_vel_w(self) -> torch.Tensor: + """Root link velocity in simulation world frame. Shape is (num_instances, 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_state_w[:, 7:13] + + @property + def root_link_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 actor frame relative to the world. + """ + return self.root_link_state_w[:, 7:10] + + @property + def root_link_ang_vel_w(self) -> torch.Tensor: + """Root link angular velocity in simulation world frame. Shape is (num_instances, 3). + + This quantity is the angular velocity of the actor frame of the root rigid body relative to the world. + """ + return self.root_link_state_w[:, 10:13] + + @property + def root_link_lin_vel_b(self) -> torch.Tensor: + """Root link linear velocity in base frame. Shape is (num_instances, 3). + + This quantity is the linear velocity of the actor frame of the root rigid body frame with respect to the + rigid body's actor frame. + """ + return math_utils.quat_rotate_inverse(self.root_link_quat_w, self.root_link_lin_vel_w) + + @property + def root_link_ang_vel_b(self) -> torch.Tensor: + """Root link angular velocity in base world frame. Shape is (num_instances, 3). + + This quantity is the angular velocity of the actor frame of the root rigid body frame with respect to the + rigid body's actor frame. + """ + return math_utils.quat_rotate_inverse(self.root_link_quat_w, self.root_link_ang_vel_w) + + @property + def root_com_pos_w(self) -> torch.Tensor: + """Root center of mass position in simulation world frame. Shape is (num_instances, 3). + + This quantity is the position of the actor frame of the root rigid body relative to the world. + """ + return self.root_com_state_w[:, :3] + + @property + def root_com_quat_w(self) -> torch.Tensor: + """Root center of mass 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 root rigid body relative to the world. + """ + return self.root_com_state_w[:, 3:7] + + @property + def root_com_vel_w(self) -> torch.Tensor: + """Root center of mass velocity in simulation world frame. Shape is (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. + """ + if self._root_link_state_w.timestamp < self._sim_timestamp: + # read data from simulation + velocity = self._root_physx_view.get_velocities() + return velocity + return self.root_com_state_w[:, 7:13] + + @property + def root_com_lin_vel_w(self) -> torch.Tensor: + """Root center of mass 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 relative to the world. + """ + if self._root_link_state_w.timestamp < self._sim_timestamp: + # read data from simulation + velocity = self._root_physx_view.get_velocities() + return velocity[:, 0:3] + return self.root_com_state_w[:, 7:10] + + @property + def root_com_ang_vel_w(self) -> torch.Tensor: + """Root center of mass angular velocity in simulation world frame. Shape is (num_instances, 3). + + This quantity is the angular velocity of the root rigid body's center of mass frame relative to the world. + """ + if self._root_link_state_w.timestamp < self._sim_timestamp: + # read data from simulation + velocity = self._root_physx_view.get_velocities() + return velocity[:, 3:6] + return self.root_com_state_w[:, 10:13] + + @property + def root_com_lin_vel_b(self) -> torch.Tensor: + """Root center of mass linear velocity in base frame. Shape is (num_instances, 3). + + This quantity is the linear velocity of the root rigid body's center of mass frame with respect to the + rigid body's actor frame. + """ + return math_utils.quat_rotate_inverse(self.root_link_quat_w, self.root_com_lin_vel_w) + + @property + def root_com_ang_vel_b(self) -> torch.Tensor: + """Root center of mass angular velocity in base world frame. Shape is (num_instances, 3). + + This quantity is the angular velocity of the root rigid body's center of mass frame with respect to the + rigid body's actor frame. + """ + return math_utils.quat_rotate_inverse(self.root_link_quat_w, self.root_com_ang_vel_w) @property def body_pos_w(self) -> torch.Tensor: @@ -278,3 +502,109 @@ def body_ang_acc_w(self) -> torch.Tensor: This quantity is the angular acceleration of the rigid bodies' center of mass frame. """ return self.body_acc_w[..., 3:6] + + # + # Link body properties + # + + @property + def body_link_pos_w(self) -> torch.Tensor: + """Positions of all bodies in simulation world frame. Shape is (num_instances, 1, 3). + + This quantity is the position of the rigid bodies' actor frame relative to the world. + """ + return self.body_link_state_w[..., :3] + + @property + def body_link_quat_w(self) -> torch.Tensor: + """Orientation (w, x, y, z) of all bodies in simulation world frame. Shape is (num_instances, 1, 4). + + This quantity is the orientation of the rigid bodies' actor frame relative to the world. + """ + return self.body_link_state_w[..., 3:7] + + @property + def body_link_vel_w(self) -> torch.Tensor: + """Velocity of all bodies in simulation world frame. Shape is (num_instances, 1, 6). + + This quantity contains the linear and angular velocities of the rigid bodies' center of mass frame + relative to the world. + """ + return self.body_link_state_w[..., 7:13] + + @property + def body_link_lin_vel_w(self) -> torch.Tensor: + """Linear velocity of all bodies in simulation world frame. Shape is (num_instances, 1, 3). + + This quantity is the linear velocity of the rigid bodies' center of mass frame relative to the world. + """ + return self.body_link_state_w[..., 7:10] + + @property + def body_link_ang_vel_w(self) -> torch.Tensor: + """Angular velocity of all bodies in simulation world frame. Shape is (num_instances, 1, 3). + + This quantity is the angular velocity of the rigid bodies' center of mass frame relative to the world. + """ + return self.body_link_state_w[..., 10:13] + + # + # Center of mass body properties + # + + @property + def body_com_pos_w(self) -> torch.Tensor: + """Positions of all bodies in simulation world frame. Shape is (num_instances, 1, 3). + + This quantity is the position of the rigid bodies' actor frame. + """ + return self.body_com_state_w[..., :3] + + @property + def body_com_quat_w(self) -> torch.Tensor: + """Orientation (w, x, y, z) of the prinicple axies of inertia of all bodies in simulation world frame. + + Shape is (num_instances, 1, 4). This quantity is the orientation of the rigid bodies' actor frame. + """ + return self.body_com_state_w[..., 3:7] + + @property + def body_com_vel_w(self) -> torch.Tensor: + """Velocity of all bodies in simulation world frame. Shape is (num_instances, 1, 6). + + This quantity contains the linear and angular velocities of the rigid bodies' center of mass frame. + """ + return self.body_com_state_w[..., 7:13] + + @property + def body_com_lin_vel_w(self) -> torch.Tensor: + """Linear velocity of all bodies in simulation world frame. Shape is (num_instances, 1, 3). + + This quantity is the linear velocity of the rigid bodies' center of mass frame. + """ + return self.body_com_state_w[..., 7:10] + + @property + def body_com_ang_vel_w(self) -> torch.Tensor: + """Angular velocity of all bodies in simulation world frame. Shape is (num_instances, 1, 3). + + This quantity is the angular velocity of the rigid bodies' center of mass frame. + """ + return self.body_com_state_w[..., 10:13] + + @property + def com_pos_b(self) -> torch.Tensor: + """Center of mass of all of the bodies in simulation world frame. Shape is (num_instances, 1, 3). + + This quantity is the center of mass location relative to its body frame. + """ + return self._root_physx_view.get_coms().to(self.device)[..., :3].view(-1, 1, 3) + + @property + def com_quat_b(self) -> torch.Tensor: + """Orientation (w,x,y,z) of the prinicple axies of inertia of all of the bodies in simulation world frame. Shape is (num_instances, 1, 4). + + This quantity is the orientation of the principles axes of inertia relative to its body frame. + """ + quat = self._root_physx_view.get_coms().to(self.device)[..., 3:7] + return math_utils.convert_quat(quat, to="wxyz").view(-1, 1, 4) diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/assets/rigid_object_collection/rigid_object_collection.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/assets/rigid_object_collection/rigid_object_collection.py index 53fca921e7..acea81ff06 100644 --- a/source/extensions/omni.isaac.lab/omni/isaac/lab/assets/rigid_object_collection/rigid_object_collection.py +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/assets/rigid_object_collection/rigid_object_collection.py @@ -105,6 +105,10 @@ def __init__(self, cfg: RigidObjectCollectionCfg): self._debug_vis_handle = None + self._root_state_dep_warn = False + self._root_pose_dep_warn = False + self._root_vel_dep_warn = False + """ Properties """ @@ -222,10 +226,56 @@ def write_object_state_to_sim( env_ids: Environment indices. If None, then all indices are used. object_ids: Object indices. If None, then all indices are used. """ + # deprecation warning + if not self._root_state_dep_warn: + omni.log.warn( + "DeprecationWarning: RigidObjectCollection.write_object_state_to_sim will be removed in a future" + " release. Please use write_object_link_state_to_sim or write_object_com_state_to_sim instead." + ) + self._root_state_dep_warn = True + # set into simulation self.write_object_pose_to_sim(object_state[..., :7], env_ids=env_ids, object_ids=object_ids) self.write_object_velocity_to_sim(object_state[..., 7:], env_ids=env_ids, object_ids=object_ids) + def write_object_com_state_to_sim( + self, + object_state: torch.Tensor, + env_ids: torch.Tensor | None = None, + object_ids: slice | torch.Tensor | None = None, + ): + """Set the object center of mass state over selected environment indices into the simulation. + The object state comprises of the cartesian position, quaternion orientation in (w, x, y, z), and linear + and angular velocity. All the quantities are in the simulation frame. + + Args: + object_state: Object state in simulation frame. Shape is (len(env_ids), len(object_ids), 13). + env_ids: Environment indices. If None, then all indices are used. + object_ids: Object indices. If None, then all indices are used. + """ + # set into simulation + self.write_object_com_pose_to_sim(object_state[..., :7], env_ids=env_ids, object_ids=object_ids) + self.write_object_com_velocity_to_sim(object_state[..., 7:], env_ids=env_ids, object_ids=object_ids) + + def write_object_link_state_to_sim( + self, + object_state: torch.Tensor, + env_ids: torch.Tensor | None = None, + object_ids: slice | torch.Tensor | None = None, + ): + """Set the object link state over selected environment indices into the simulation. + The object state comprises of the cartesian position, quaternion orientation in (w, x, y, z), and linear + and angular velocity. All the quantities are in the simulation frame. + + Args: + object_state: Object state in simulation frame. Shape is (len(env_ids), len(object_ids), 13). + env_ids: Environment indices. If None, then all indices are used. + object_ids: Object indices. If None, then all indices are used. + """ + # set into simulation + self.write_object_link_pose_to_sim(object_state[..., :7], env_ids=env_ids, object_ids=object_ids) + self.write_object_link_velocity_to_sim(object_state[..., 7:], env_ids=env_ids, object_ids=object_ids) + def write_object_pose_to_sim( self, object_pose: torch.Tensor, @@ -241,6 +291,32 @@ def write_object_pose_to_sim( env_ids: Environment indices. If None, then all indices are used. object_ids: Object indices. If None, then all indices are used. """ + # deprecation warning + if not self._root_pose_dep_warn: + omni.log.warn( + "DeprecationWarning: RigidObjectCollection.write_object_pose_to_sim will be removed in a future" + " release. Please use write_object_link_pose_to_sim or write_object_com_pose_to_sim instead." + ) + self._root_pose_dep_warn = True + + self.write_object_link_pose_to_sim(object_pose, env_ids, object_ids) + + def write_object_link_pose_to_sim( + self, + object_pose: torch.Tensor, + env_ids: torch.Tensor | None = None, + object_ids: slice | torch.Tensor | None = None, + ): + """Set the object pose over selected environment and object indices into the simulation. + + The object pose comprises of the cartesian position and quaternion orientation in (w, x, y, z). + + Args: + object_pose: Object poses in simulation frame. Shape is (len(env_ids), len(object_ids), 7). + env_ids: Environment indices. If None, then all indices are used. + object_ids: Object indices. If None, then all indices are used. + """ + # resolve all indices # -- env_ids if env_ids is None: @@ -250,14 +326,55 @@ def write_object_pose_to_sim( object_ids = self._ALL_OBJ_INDICES # note: we need to do this here since tensors are not set into simulation until step. # set into internal buffers + self._data.object_link_state_w[env_ids[:, None], object_ids, :7] = object_pose.clone() + self._data._ignore_dep_warn = True self._data.object_state_w[env_ids[:, None], object_ids, :7] = object_pose.clone() + self._data._ignore_dep_warn = False # convert the quaternion from wxyz to xyzw - poses_xyzw = self._data.object_state_w[..., :7].clone() + poses_xyzw = self._data.object_link_state_w[..., :7].clone() poses_xyzw[..., 3:] = math_utils.convert_quat(poses_xyzw[..., 3:], to="xyzw") # set into simulation view_ids = self._env_obj_ids_to_view_ids(env_ids, object_ids) self.root_physx_view.set_transforms(self.reshape_data_to_view(poses_xyzw), indices=view_ids) + def write_object_com_pose_to_sim( + self, + object_pose: torch.Tensor, + env_ids: torch.Tensor | None = None, + object_ids: slice | torch.Tensor | None = None, + ): + """Set the object center of mass pose over selected environment indices into the simulation. + The object pose comprises of the cartesian position and quaternion orientation in (w, x, y, z). + The orientation is the orientation of the principle axes of inertia. + Args: + object_pose: Object poses in simulation frame. Shape is (len(env_ids), len(object_ids), 7). + env_ids: Environment indices. If None, then all indices are used. + object_ids: Object indices. If None, then all indices are used. + """ + + # resolve all indices + if env_ids is None: + local_env_ids = slice(env_ids) + else: + local_env_ids = env_ids + if object_ids is None: + local_object_ids = slice(object_ids) + else: + local_object_ids = object_ids + + com_pos = self.data.com_pos_b[local_env_ids, local_object_ids, :] + com_quat = self.data.com_quat_b[local_env_ids, local_object_ids, :] + + object_link_pos, object_link_quat = math_utils.combine_frame_transforms( + object_pose[..., :3], + object_pose[..., 3:7], + math_utils.quat_rotate(math_utils.quat_inv(com_quat), -com_pos), + math_utils.quat_inv(com_quat), + ) + + object_link_pose = torch.cat((object_link_pos, object_link_quat), dim=-1) + self.write_object_link_pose_to_sim(object_pose=object_link_pose, env_ids=env_ids, object_ids=object_ids) + def write_object_velocity_to_sim( self, object_velocity: torch.Tensor, @@ -266,6 +383,29 @@ def write_object_velocity_to_sim( ): """Set the object velocity over selected environment and object indices into the simulation. + Args: + object_velocity: Object velocities in simulation frame. Shape is (len(env_ids), len(object_ids), 6). + env_ids: Environment indices. If None, then all indices are used. + object_ids: Object indices. If None, then all indices are used. + """ + # deprecation warning + if not self._root_vel_dep_warn: + omni.log.warn( + "DeprecationWarning: RigidObjectCollection.write_object_velocity_to_sim will be removed in a future" + " release. Please use write_object_link_velocity_to_sim or write_object_com_velocity_to_sim instead." + ) + self._root_vel_dep_warn = True + + self.write_object_com_velocity_to_sim(object_velocity=object_velocity, env_ids=env_ids, object_ids=object_ids) + + def write_object_com_velocity_to_sim( + self, + object_velocity: torch.Tensor, + env_ids: torch.Tensor | None = None, + object_ids: slice | torch.Tensor | None = None, + ): + """Set the object center of mass velocity over selected environment and object indices into the simulation. + Args: object_velocity: Object velocities in simulation frame. Shape is (len(env_ids), len(object_ids), 6). env_ids: Environment indices. If None, then all indices are used. @@ -279,13 +419,52 @@ def write_object_velocity_to_sim( if object_ids is None: object_ids = self._ALL_OBJ_INDICES + self._data.object_com_state_w[env_ids[:, None], object_ids, 7:] = object_velocity.clone() + self._data._ignore_dep_warn = True self._data.object_state_w[env_ids[:, None], object_ids, 7:] = object_velocity.clone() + self._data._ignore_dep_warn = False self._data.object_acc_w[env_ids[:, None], object_ids] = 0.0 # set into simulation view_ids = self._env_obj_ids_to_view_ids(env_ids, object_ids) self.root_physx_view.set_velocities( - self.reshape_data_to_view(self._data.object_state_w[..., 7:]), indices=view_ids + self.reshape_data_to_view(self._data.object_com_state_w[..., 7:]), indices=view_ids + ) + + def write_object_link_velocity_to_sim( + self, + object_velocity: torch.Tensor, + env_ids: torch.Tensor | None = None, + object_ids: slice | torch.Tensor | None = None, + ): + """Set the object 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 object's frame rather than the objects center of mass. + Args: + object_velocity: Object velocities in simulation frame. Shape is (len(env_ids), len(object_ids), 6). + env_ids: Environment indices. If None, then all indices are used. + object_ids: Object indices. If None, then all indices are used. + """ + # resolve all indices + if env_ids is None: + local_env_ids = slice(env_ids) + else: + local_env_ids = env_ids + if object_ids is None: + local_object_ids = slice(object_ids) + else: + local_object_ids = object_ids + + object_com_velocity = object_velocity.clone() + quat = self.data.object_link_state_w[local_env_ids, local_object_ids, 3:7] + com_pos_b = self.data.com_pos_b[local_env_ids, local_object_ids, :] + # transform given velocity to center of mass + object_com_velocity[..., :3] += torch.linalg.cross( + object_com_velocity[..., 3:], math_utils.quat_rotate(quat, com_pos_b), dim=-1 + ) + # write center of mass velocity to sim + self.write_object_com_velocity_to_sim( + object_velocity=object_com_velocity, env_ids=env_ids, object_ids=object_ids ) """ diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/assets/rigid_object_collection/rigid_object_collection_data.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/assets/rigid_object_collection/rigid_object_collection_data.py index c9930cf740..accae3befc 100644 --- a/source/extensions/omni.isaac.lab/omni/isaac/lab/assets/rigid_object_collection/rigid_object_collection_data.py +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/assets/rigid_object_collection/rigid_object_collection_data.py @@ -6,6 +6,7 @@ import torch import weakref +import omni.log import omni.physics.tensors.impl.api as physx import omni.isaac.lab.utils.math as math_utils @@ -69,8 +70,14 @@ def __init__(self, root_physx_view: physx.RigidBodyView, num_objects: int, devic # Initialize the lazy buffers. self._object_state_w = TimestampedBuffer() + self._object_link_state_w = TimestampedBuffer() + self._object_com_state_w = TimestampedBuffer() self._object_acc_w = TimestampedBuffer() + # deprecation warning check + self._root_state_dep_warn = False + self._ignore_dep_warn = False + def update(self, dt: float): """Updates the data for the rigid object collection. @@ -119,6 +126,14 @@ def object_state_w(self): The position and orientation are of the rigid body's actor frame. Meanwhile, the linear and angular velocities are of the rigid body's center of mass frame. """ + + if not self._root_state_dep_warn and not self._ignore_dep_warn: + omni.log.warn( + "DeprecationWarning: object_state_w and it's derived properties will be deprecated in a future release." + " Please use object_link_state_w or object_com_state_w." + ) + self._root_state_dep_warn = True + if self._object_state_w.timestamp < self._sim_timestamp: # read data from simulation pose = self._reshape_view_to_data(self._root_physx_view.get_transforms().clone()) @@ -129,6 +144,53 @@ def object_state_w(self): self._object_state_w.timestamp = self._sim_timestamp return self._object_state_w.data + @property + def object_link_state_w(self): + """Object center of mass state ``[pos, quat, lin_vel, ang_vel]`` in simulation world frame. Shape is (num_instances, num_objects, 13). + + The position, quaternion, and linear/angular velocity are of the rigid body root frame relative to the + world. + """ + if self._object_link_state_w.timestamp < self._sim_timestamp: + # read data from simulation + pose = self._reshape_view_to_data(self._root_physx_view.get_transforms().clone()) + pose[..., 3:7] = math_utils.convert_quat(pose[..., 3:7], to="wxyz") + velocity = self._reshape_view_to_data(self._root_physx_view.get_velocities()) + + # adjust linear velocity to link from center of mass + velocity[..., :3] += torch.linalg.cross( + velocity[..., 3:], math_utils.quat_rotate(pose[..., 3:7], -self.com_pos_b[..., :]), dim=-1 + ) + + # set the buffer data and timestamp + self._object_link_state_w.data = torch.cat((pose, velocity), dim=-1) + self._object_link_state_w.timestamp = self._sim_timestamp + return self._object_link_state_w.data + + @property + def object_com_state_w(self): + """Object state ``[pos, quat, lin_vel, ang_vel]`` in simulation world frame. Shape is (num_instances, num_objects, 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 the orientation principle axes of inertia. + """ + + if self._object_com_state_w.timestamp < self._sim_timestamp: + # read data from simulation + pose = self._reshape_view_to_data(self._root_physx_view.get_transforms().clone()) + pose[..., 3:7] = math_utils.convert_quat(pose[..., 3:7], to="wxyz") + velocity = self._reshape_view_to_data(self._root_physx_view.get_velocities()) + + # adjust pose to center of mass + pos, quat = math_utils.combine_frame_transforms( + pose[..., :3], pose[..., 3:7], self.com_pos_b[..., :], self.com_quat_b[..., :] + ) + + # set the buffer data and timestamp + self._object_com_state_w.data = torch.cat((pos, quat, velocity), dim=-1) + self._object_com_state_w.timestamp = self._sim_timestamp + return self._object_com_state_w.data + @property def object_acc_w(self): """Acceleration of all objects. Shape is (num_instances, num_objects, 6). @@ -144,7 +206,7 @@ def object_acc_w(self): @property def projected_gravity_b(self): """Projection of the gravity direction on base frame. Shape is (num_instances, num_objects, 3).""" - return math_utils.quat_rotate_inverse(self.object_quat_w, self.GRAVITY_VEC_W) + return math_utils.quat_rotate_inverse(self.object_link_quat_w, self.GRAVITY_VEC_W) @property def heading_w(self): @@ -154,7 +216,7 @@ def heading_w(self): This quantity is computed by assuming that the forward-direction of the base frame is along x-direction, i.e. :math:`(1, 0, 0)`. """ - forward_w = math_utils.quat_apply(self.object_quat_w, self.FORWARD_VEC_B) + forward_w = math_utils.quat_apply(self.object_link_quat_w, self.FORWARD_VEC_B) return torch.atan2(forward_w[..., 1], forward_w[..., 0]) ## @@ -235,6 +297,162 @@ def object_ang_acc_w(self) -> torch.Tensor: """ return self.object_acc_w[..., 3:6] + @property + def object_link_pos_w(self) -> torch.Tensor: + """Object link position in simulation world frame. Shape is (num_instances, num_objects, 3). + + This quantity is the position of the actor frame of the rigid bodies. + """ + if self._object_state_w.timestamp < self._sim_timestamp: + # read data from simulation + pose = self._reshape_view_to_data(self._root_physx_view.get_transforms().clone()) + return pose[..., :3] + return self.object_link_state_w[..., :3] + + @property + def object_link_quat_w(self) -> torch.Tensor: + """Object link orientation (w, x, y, z) in simulation world frame. Shape is (num_instances, num_objects, 4). + + This quantity is the orientation of the actor frame of the rigid bodies. + """ + if self._object_state_w.timestamp < self._sim_timestamp: + # read data from simulation + pose = self._reshape_view_to_data(self._root_physx_view.get_transforms().clone()) + pose[..., 3:7] = math_utils.convert_quat(pose[..., 3:7], to="wxyz") + return pose[..., 3:7] + return self.object_link_state_w[..., 3:7] + + @property + def object_link_vel_w(self) -> torch.Tensor: + """Object link velocity in simulation world frame. Shape is (num_instances, num_objects, 6). + + This quantity contains the linear and angular velocities of the rigid bodies' actor frame. + """ + return self.object_link_state_w[..., 7:13] + + @property + def object_link_lin_vel_w(self) -> torch.Tensor: + """Object link linear velocity in simulation world frame. Shape is (num_instances, num_objects, 3). + + This quantity is the linear velocity of the rigid bodies' actor frame. + """ + return self.object_link_state_w[..., 7:10] + + @property + def object_link_ang_vel_w(self) -> torch.Tensor: + """Object link angular velocity in simulation world frame. Shape is (num_instances, num_objects, 3). + + This quantity is the angular velocity of the rigid bodies' actor frame. + """ + return self.object_link_state_w[..., 10:13] + + @property + def object_link_lin_vel_b(self) -> torch.Tensor: + """Object link linear velocity in base frame. Shape is (num_instances, num_objects, 3). + + This quantity is the linear velocity of the actor frame of the root rigid body frame with respect to the + rigid body's actor frame. + """ + return math_utils.quat_rotate_inverse(self.object_link_quat_w, self.object_link_lin_vel_w) + + @property + def object_link_ang_vel_b(self) -> torch.Tensor: + """Object link angular velocity in base world frame. Shape is (num_instances, num_objects, 3). + + This quantity is the angular velocity of the actor frame of the root rigid body frame with respect to the + rigid body's actor frame. + """ + return math_utils.quat_rotate_inverse(self.object_link_quat_w, self.object_link_ang_vel_w) + + @property + def object_com_pos_w(self) -> torch.Tensor: + """Object center of mass position in simulation world frame. Shape is (num_instances, num_objects, 3). + + This quantity is the position of the center of mass frame of the rigid bodies. + """ + return self.object_com_state_w[..., :3] + + @property + def object_com_quat_w(self) -> torch.Tensor: + """Object center of mass orientation (w, x, y, z) in simulation world frame. Shape is (num_instances, num_objects, 4). + + This quantity is the orientation of the center of mass frame of the rigid bodies. + """ + return self.object_com_state_w[..., 3:7] + + @property + def object_com_vel_w(self) -> torch.Tensor: + """Object center of mass velocity in simulation world frame. Shape is (num_instances, num_objects, 6). + + This quantity contains the linear and angular velocities of the rigid bodies' center of mass frame. + """ + if self._object_state_w.timestamp < self._sim_timestamp: + # read data from simulation + velocity = self._reshape_view_to_data(self._root_physx_view.get_velocities()) + return velocity + return self.object_com_state_w[..., 7:13] + + @property + def object_com_lin_vel_w(self) -> torch.Tensor: + """Object center of mass linear velocity in simulation world frame. Shape is (num_instances, num_objects, 3). + + This quantity is the linear velocity of the rigid bodies' center of mass frame. + """ + if self._object_state_w.timestamp < self._sim_timestamp: + # read data from simulation + velocity = self._reshape_view_to_data(self._root_physx_view.get_velocities()) + return velocity[..., 0:3] + return self.object_com_state_w[..., 7:10] + + @property + def object_com_ang_vel_w(self) -> torch.Tensor: + """Object center of mass angular velocity in simulation world frame. Shape is (num_instances, num_objects, 3). + + This quantity is the angular velocity of the rigid bodies' center of mass frame. + """ + if self._object_state_w.timestamp < self._sim_timestamp: + # read data from simulation + velocity = self._reshape_view_to_data(self._root_physx_view.get_velocities()) + return velocity[..., 3:6] + return self.object_com_state_w[..., 10:13] + + @property + def object_com_lin_vel_b(self) -> torch.Tensor: + """Object center of mass linear velocity in base frame. Shape is (num_instances, num_objects, 3). + + This quantity is the linear velocity of the center of mass frame of the root rigid body frame with respect to the + rigid body's actor frame. + """ + return math_utils.quat_rotate_inverse(self.object_link_quat_w, self.object_com_lin_vel_w) + + @property + def object_com_ang_vel_b(self) -> torch.Tensor: + """Object center of mass angular velocity in base world frame. Shape is (num_instances, num_objects, 3). + + This quantity is the angular velocity of the center of mass frame of the root rigid body frame with respect to the + rigid body's actor frame. + """ + return math_utils.quat_rotate_inverse(self.object_link_quat_w, self.object_com_ang_vel_w) + + @property + def com_pos_b(self) -> torch.Tensor: + """Center of mass of all of the bodies in simulation world frame. Shape is (num_instances, 1, 3). + + This quantity is the center of mass location relative to its body frame. + """ + pos = self._root_physx_view.get_coms().to(self.device)[..., :3] + return self._reshape_view_to_data(pos) + + @property + def com_quat_b(self) -> torch.Tensor: + """Orientation (w,x,y,z) of the prinicple axies of inertia of all of the bodies in simulation world frame. Shape is (num_instances, 1, 4). + + This quantity is the orientation of the principles axes of inertia relative to its body frame. + """ + quat = self._root_physx_view.get_coms().to(self.device)[..., 3:7].view(self.num_instances, self.num_objects, 4) + quat_wxyz = math_utils.convert_quat(quat, to="wxyz") + return self._reshape_view_to_data(quat_wxyz) + ## # Helpers. ## diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/mdp/actions/non_holonomic_actions.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/mdp/actions/non_holonomic_actions.py index b95a058339..4d002810e5 100644 --- a/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/mdp/actions/non_holonomic_actions.py +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/mdp/actions/non_holonomic_actions.py @@ -150,7 +150,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_link_quat_w[:, 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/extensions/omni.isaac.lab/omni/isaac/lab/envs/mdp/actions/task_space_actions.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/mdp/actions/task_space_actions.py index 5bd38957be..8a96a2523c 100644 --- a/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/mdp/actions/task_space_actions.py +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/mdp/actions/task_space_actions.py @@ -142,7 +142,7 @@ def jacobian_w(self) -> torch.Tensor: @property def jacobian_b(self) -> torch.Tensor: jacobian = self.jacobian_w - base_rot = self._asset.data.root_quat_w + base_rot = self._asset.data.root_link_quat_w 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:, :]) @@ -192,12 +192,12 @@ 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_pose_w = self._asset.data.body_state_w[:, self._body_idx, :7] - root_pose_w = self._asset.data.root_state_w[:, :7] + ee_pos_w = self._asset.data.body_link_pos_w[:, self._body_idx] + ee_quat_w = self._asset.data.body_link_quat_w[:, self._body_idx] + root_pos_w = self._asset.data.root_link_pos_w + root_quat_w = self._asset.data.root_link_quat_w # compute the pose of the body in the root frame - ee_pose_b, ee_quat_b = math_utils.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, ee_quat_b = math_utils.subtract_frame_transforms(root_pos_w, root_quat_w, ee_pos_w, ee_quat_w) # account for the offset if self.cfg.body_offset is not None: ee_pose_b, ee_quat_b = math_utils.combine_frame_transforms( @@ -395,7 +395,7 @@ def jacobian_w(self) -> torch.Tensor: @property def jacobian_b(self) -> torch.Tensor: jacobian = self.jacobian_w - base_rot = self._asset.data.root_quat_w + base_rot = self._asset.data.root_link_quat_w 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:, :]) @@ -556,10 +556,14 @@ 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[:] = self._asset.data.body_state_w[:, self._ee_body_idx, :7] + self._ee_pose_w[:, 0:3] = self._asset.data.body_link_pos_w[:, self._ee_body_idx] + self._ee_pose_w[:, 3:7] = self._asset.data.body_link_quat_w[:, 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( - self._asset.data.root_pos_w, self._asset.data.root_quat_w, self._ee_pose_w[:, 0:3], self._ee_pose_w[:, 3:7] + self._asset.data.root_link_pos_w, + self._asset.data.root_link_quat_w, + self._ee_pose_w[:, 0:3], + self._ee_pose_w[:, 3:7], ) # Account for the offset if self.cfg.body_offset is not None: @@ -572,13 +576,17 @@ 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[:] = self._asset.data.body_vel_w[:, self._ee_body_idx, :] + self._ee_vel_w[:] = self._asset.data.body_com_vel_w[:, self._ee_body_idx, :] # Compute the relative velocity in the world frame - relative_vel_w = self._ee_vel_w - self._asset.data.root_vel_w + relative_vel_w = self._ee_vel_w - self._asset.data.root_com_vel_w # Convert ee velocities from world to root frame - self._ee_vel_b[:, 0:3] = math_utils.quat_rotate_inverse(self._asset.data.root_quat_w, relative_vel_w[:, 0:3]) - self._ee_vel_b[:, 3:6] = math_utils.quat_rotate_inverse(self._asset.data.root_quat_w, relative_vel_w[:, 3:6]) + self._ee_vel_b[:, 0:3] = math_utils.quat_rotate_inverse( + self._asset.data.root_link_quat_w, relative_vel_w[:, 0:3] + ) + self._ee_vel_b[:, 3:6] = math_utils.quat_rotate_inverse( + self._asset.data.root_link_quat_w, relative_vel_w[:, 3:6] + ) # Account for the offset if self.cfg.body_offset is not None: @@ -595,7 +603,7 @@ def _compute_ee_force(self): self._contact_sensor.update(self._sim_dt) self._ee_force_w[:] = self._contact_sensor.data.net_forces_w[:, 0, :] # type: ignore # Rotate forces and torques into root frame - self._ee_force_b[:] = math_utils.quat_rotate_inverse(self._asset.data.root_quat_w, self._ee_force_w) + self._ee_force_b[:] = math_utils.quat_rotate_inverse(self._asset.data.root_link_quat_w, self._ee_force_w) def _compute_task_frame_pose(self): """Computes the pose of the task frame in root frame.""" @@ -604,8 +612,8 @@ 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( - self._asset.data.root_pos_w, - self._asset.data.root_quat_w, + self._asset.data.root_link_pos_w, + self._asset.data.root_link_quat_w, self._task_frame_transformer.data.target_pos_w[:, 0, :], self._task_frame_transformer.data.target_quat_w[:, 0, :], ) diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/mdp/commands/pose_2d_command.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/mdp/commands/pose_2d_command.py index f933760e07..f38f5fbd6c 100644 --- a/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/mdp/commands/pose_2d_command.py +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/mdp/commands/pose_2d_command.py @@ -81,7 +81,9 @@ def command(self) -> torch.Tensor: def _update_metrics(self): # logs data - self.metrics["error_pos_2d"] = torch.norm(self.pos_command_w[:, :2] - self.robot.data.root_pos_w[:, :2], dim=1) + self.metrics["error_pos_2d"] = torch.norm( + self.pos_command_w[:, :2] - self.robot.data.root_link_pos_w[:, :2], dim=1 + ) self.metrics["error_heading"] = torch.abs(wrap_to_pi(self.heading_command_w - self.robot.data.heading_w)) def _resample_command(self, env_ids: Sequence[int]): @@ -95,7 +97,7 @@ def _resample_command(self, env_ids: Sequence[int]): if self.cfg.simple_heading: # set heading command to point towards target - target_vec = self.pos_command_w[env_ids] - self.robot.data.root_pos_w[env_ids] + target_vec = self.pos_command_w[env_ids] - self.robot.data.root_link_pos_w[env_ids] target_direction = torch.atan2(target_vec[:, 1], target_vec[:, 0]) flipped_target_direction = wrap_to_pi(target_direction + torch.pi) @@ -116,8 +118,8 @@ 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 - self.robot.data.root_pos_w[:, :3] - self.pos_command_b[:] = quat_rotate_inverse(yaw_quat(self.robot.data.root_quat_w), target_vec) + target_vec = self.pos_command_w - self.robot.data.root_link_pos_w[:, :3] + self.pos_command_b[:] = quat_rotate_inverse(yaw_quat(self.robot.data.root_link_quat_w), target_vec) self.heading_command_b[:] = wrap_to_pi(self.heading_command_w - self.robot.data.heading_w) def _set_debug_vis_impl(self, debug_vis: bool): @@ -182,7 +184,7 @@ def _resample_command(self, env_ids: Sequence[int]): if self.cfg.simple_heading: # set heading command to point towards target - target_vec = self.pos_command_w[env_ids] - self.robot.data.root_pos_w[env_ids] + target_vec = self.pos_command_w[env_ids] - self.robot.data.root_link_pos_w[env_ids] target_direction = torch.atan2(target_vec[:, 1], target_vec[:, 0]) flipped_target_direction = wrap_to_pi(target_direction + torch.pi) diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/mdp/commands/pose_command.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/mdp/commands/pose_command.py index 328834540e..4afb50d5ad 100644 --- a/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/mdp/commands/pose_command.py +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/mdp/commands/pose_command.py @@ -92,8 +92,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( - self.robot.data.root_pos_w, - self.robot.data.root_quat_w, + self.robot.data.root_link_pos_w, + self.robot.data.root_link_quat_w, self.pose_command_b[:, :3], self.pose_command_b[:, 3:], ) @@ -101,8 +101,8 @@ def _update_metrics(self): 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_link_state_w[:, self.body_idx, :3], + self.robot.data.body_link_state_w[:, self.body_idx, 3:7], ) self.metrics["position_error"] = torch.norm(pos_error, dim=-1) self.metrics["orientation_error"] = torch.norm(rot_error, dim=-1) @@ -151,5 +151,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_pose_w = self.robot.data.body_state_w[:, self.body_idx] - self.current_pose_visualizer.visualize(body_pose_w[:, :3], body_pose_w[:, 3:7]) + body_link_state_w = self.robot.data.body_link_state_w[:, self.body_idx] + self.current_pose_visualizer.visualize(body_link_state_w[:, :3], body_link_state_w[:, 3:7]) diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/mdp/commands/velocity_command.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/mdp/commands/velocity_command.py index 2cabd86dba..93a4f51d10 100644 --- a/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/mdp/commands/velocity_command.py +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/mdp/commands/velocity_command.py @@ -114,10 +114,10 @@ def _update_metrics(self): max_command_step = max_command_time / self._env.step_dt # logs data self.metrics["error_vel_xy"] += ( - torch.norm(self.vel_command_b[:, :2] - self.robot.data.root_lin_vel_b[:, :2], dim=-1) / max_command_step + torch.norm(self.vel_command_b[:, :2] - self.robot.data.root_com_lin_vel_b[:, :2], dim=-1) / max_command_step ) self.metrics["error_vel_yaw"] += ( - torch.abs(self.vel_command_b[:, 2] - self.robot.data.root_ang_vel_b[:, 2]) / max_command_step + torch.abs(self.vel_command_b[:, 2] - self.robot.data.root_com_ang_vel_b[:, 2]) / max_command_step ) def _resample_command(self, env_ids: Sequence[int]): @@ -184,11 +184,11 @@ def _debug_vis_callback(self, event): return # get marker location # -- base state - base_pos_w = self.robot.data.root_pos_w.clone() + base_pos_w = self.robot.data.root_link_pos_w.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(self.robot.data.root_lin_vel_b[:, :2]) + vel_arrow_scale, vel_arrow_quat = self._resolve_xy_velocity_to_arrow(self.robot.data.root_com_lin_vel_b[:, :2]) # display markers self.goal_vel_visualizer.visualize(base_pos_w, vel_des_arrow_quat, vel_des_arrow_scale) self.current_vel_visualizer.visualize(base_pos_w, vel_arrow_quat, vel_arrow_scale) @@ -209,7 +209,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 = self.robot.data.root_quat_w + base_quat_w = self.robot.data.root_link_quat_w arrow_quat = math_utils.quat_mul(base_quat_w, arrow_quat) return arrow_scale, arrow_quat diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/mdp/events.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/mdp/events.py index bd7532a14a..44a58374e1 100644 --- a/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/mdp/events.py +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/mdp/events.py @@ -624,13 +624,13 @@ def push_by_setting_velocity( asset: RigidObject | Articulation = env.scene[asset_cfg.name] # velocities - vel_w = asset.data.root_vel_w[env_ids] + vel_w = asset.data.root_com_vel_w[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) vel_w[:] = math_utils.sample_uniform(ranges[:, 0], ranges[:, 1], vel_w.shape, device=asset.device) # set the velocities into the physics simulation - asset.write_root_velocity_to_sim(vel_w, env_ids=env_ids) + asset.write_root_com_velocity_to_sim(vel_w, env_ids=env_ids) def reset_root_state_uniform( @@ -674,8 +674,8 @@ def reset_root_state_uniform( velocities = root_states[:, 7:13] + rand_samples # set into the physics simulation - asset.write_root_pose_to_sim(torch.cat([positions, orientations], dim=-1), env_ids=env_ids) - asset.write_root_velocity_to_sim(velocities, env_ids=env_ids) + asset.write_root_link_pose_to_sim(torch.cat([positions, orientations], dim=-1), env_ids=env_ids) + asset.write_root_com_velocity_to_sim(velocities, env_ids=env_ids) def reset_root_state_with_random_orientation( @@ -726,8 +726,8 @@ def reset_root_state_with_random_orientation( velocities = root_states[:, 7:13] + rand_samples # set into the physics simulation - asset.write_root_pose_to_sim(torch.cat([positions, orientations], dim=-1), env_ids=env_ids) - asset.write_root_velocity_to_sim(velocities, env_ids=env_ids) + asset.write_root_link_pose_to_sim(torch.cat([positions, orientations], dim=-1), env_ids=env_ids) + asset.write_root_com_velocity_to_sim(velocities, env_ids=env_ids) def reset_root_state_from_terrain( @@ -793,8 +793,8 @@ def reset_root_state_from_terrain( velocities = asset.data.default_root_state[:, 7:13] + rand_samples # set into the physics simulation - asset.write_root_pose_to_sim(torch.cat([positions, orientations], dim=-1), env_ids=env_ids) - asset.write_root_velocity_to_sim(velocities, env_ids=env_ids) + asset.write_root_link_pose_to_sim(torch.cat([positions, orientations], dim=-1), env_ids=env_ids) + asset.write_root_com_velocity_to_sim(velocities, env_ids=env_ids) def reset_joints_by_scale( @@ -914,14 +914,16 @@ def reset_scene_to_default(env: ManagerBasedEnv, env_ids: torch.Tensor): default_root_state = rigid_object.data.default_root_state[env_ids].clone() default_root_state[:, 0:3] += env.scene.env_origins[env_ids] # set into the physics simulation - rigid_object.write_root_state_to_sim(default_root_state, env_ids=env_ids) + rigid_object.write_root_link_pose_to_sim(default_root_state[:, :7], env_ids=env_ids) + rigid_object.write_root_com_velocity_to_sim(default_root_state[:, 7:], env_ids=env_ids) # articulations for articulation_asset in env.scene.articulations.values(): # obtain default and deal with the offset for env origins default_root_state = articulation_asset.data.default_root_state[env_ids].clone() default_root_state[:, 0:3] += env.scene.env_origins[env_ids] # set into the physics simulation - articulation_asset.write_root_state_to_sim(default_root_state, env_ids=env_ids) + articulation_asset.write_root_link_pose_to_sim(default_root_state[:, :7], env_ids=env_ids) + articulation_asset.write_root_com_velocity_to_sim(default_root_state[:, 7:], env_ids=env_ids) # obtain default joint positions default_joint_pos = articulation_asset.data.default_joint_pos[env_ids].clone() default_joint_vel = articulation_asset.data.default_joint_vel[env_ids].clone() diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/mdp/observations.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/mdp/observations.py index 2c97cf151e..8df6dbbead 100644 --- a/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/mdp/observations.py +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/mdp/observations.py @@ -34,21 +34,21 @@ 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 asset.data.root_pos_w[:, 2].unsqueeze(-1) + return asset.data.root_link_pos_w[:, 2].unsqueeze(-1) def base_lin_vel(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor: """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 asset.data.root_lin_vel_b + return asset.data.root_com_lin_vel_b def base_ang_vel(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor: """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 asset.data.root_ang_vel_b + return asset.data.root_com_ang_vel_b def projected_gravity(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor: @@ -62,7 +62,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 asset.data.root_pos_w - env.scene.env_origins + return asset.data.root_link_pos_w - env.scene.env_origins def root_quat_w( @@ -77,7 +77,7 @@ def root_quat_w( # extract the used quantities (to enable type-hinting) asset: RigidObject = env.scene[asset_cfg.name] - quat = asset.data.root_quat_w + quat = asset.data.root_link_quat_w # make the quaternion real-part positive if configured return math_utils.quat_unique(quat) if make_quat_unique else quat @@ -86,14 +86,14 @@ 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 asset.data.root_lin_vel_w + return asset.data.root_com_lin_vel_w def root_ang_vel_w(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor: """Asset root angular velocity in the environment frame.""" # extract the used quantities (to enable type-hinting) asset: RigidObject = env.scene[asset_cfg.name] - return asset.data.root_ang_vel_w + return asset.data.root_com_ang_vel_w """ diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/mdp/rewards.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/mdp/rewards.py index 9b3fe7440c..cd5564a08b 100644 --- a/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/mdp/rewards.py +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/mdp/rewards.py @@ -77,14 +77,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(asset.data.root_lin_vel_b[:, 2]) + return torch.square(asset.data.root_com_lin_vel_b[:, 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(asset.data.root_ang_vel_b[:, :2]), dim=1) + return torch.sum(torch.square(asset.data.root_com_ang_vel_b[:, :2]), dim=1) def flat_orientation_l2(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor: @@ -119,7 +119,7 @@ def base_height_l2( # Use the provided target height directly for flat terrain adjusted_target_height = target_height # Compute the L2 squared penalty - return torch.square(asset.data.root_pos_w[:, 2] - adjusted_target_height) + return torch.square(asset.data.root_link_pos_w[:, 2] - adjusted_target_height) def body_lin_acc_l2(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor: @@ -292,7 +292,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] - asset.data.root_lin_vel_b[:, :2]), + torch.square(env.command_manager.get_command(command_name)[:, :2] - asset.data.root_com_lin_vel_b[:, :2]), dim=1, ) return torch.exp(-lin_vel_error / std**2) @@ -305,5 +305,7 @@ def track_ang_vel_z_exp( # extract the used quantities (to enable type-hinting) asset: RigidObject = env.scene[asset_cfg.name] # compute the error - ang_vel_error = torch.square(env.command_manager.get_command(command_name)[:, 2] - asset.data.root_ang_vel_b[:, 2]) + ang_vel_error = torch.square( + env.command_manager.get_command(command_name)[:, 2] - asset.data.root_com_ang_vel_b[:, 2] + ) return torch.exp(-ang_vel_error / std**2) diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/mdp/terminations.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/mdp/terminations.py index a91bb39581..8f4017597b 100644 --- a/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/mdp/terminations.py +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/mdp/terminations.py @@ -69,7 +69,7 @@ def root_height_below_minimum( """ # extract the used quantities (to enable type-hinting) asset: RigidObject = env.scene[asset_cfg.name] - return asset.data.root_pos_w[:, 2] < minimum_height + return asset.data.root_link_pos_w[:, 2] < minimum_height """ diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/ui/viewport_camera_controller.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/ui/viewport_camera_controller.py index a1fef9f445..bf7463c9ba 100644 --- a/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/ui/viewport_camera_controller.py +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/ui/viewport_camera_controller.py @@ -156,7 +156,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 - self.viewer_origin = self._env.scene[self.cfg.asset_name].data.root_pos_w[self.cfg.env_index] + self.viewer_origin = self._env.scene[self.cfg.asset_name].data.root_link_pos_w[self.cfg.env_index] # update the camera view self.update_view_location() diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/scene/interactive_scene.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/scene/interactive_scene.py index 08be6ed979..25dd4bcdb0 100644 --- a/source/extensions/omni.isaac.lab/omni/isaac/lab/scene/interactive_scene.py +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/scene/interactive_scene.py @@ -364,10 +364,10 @@ 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"] = articulation.data.root_state_w[:, :7].clone() + asset_state["root_pose"] = articulation.data.root_link_state_w[:, :7].clone() if is_relative: asset_state["root_pose"][:, :3] -= self.env_origins - asset_state["root_velocity"] = articulation.data.root_vel_w.clone() + asset_state["root_velocity"] = articulation.data.root_com_vel_w.clone() asset_state["joint_position"] = articulation.data.joint_pos.clone() asset_state["joint_velocity"] = articulation.data.joint_vel.clone() state["articulation"][asset_name] = asset_state @@ -384,10 +384,10 @@ def get_state(self, is_relative: bool = False) -> dict[str, dict[str, dict[str, 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_link_state_w[:, :7].clone() if is_relative: asset_state["root_pose"][:, :3] -= self.env_origins - asset_state["root_velocity"] = rigid_object.data.root_vel_w.clone() + asset_state["root_velocity"] = rigid_object.data.root_com_vel_w.clone() state["rigid_object"][asset_name] = asset_state return state @@ -439,8 +439,8 @@ def reset_to( if is_relative: root_pose[:, :3] += self.env_origins[env_ids] root_velocity = asset_state["root_velocity"].clone() - articulation.write_root_pose_to_sim(root_pose, env_ids=env_ids) - articulation.write_root_velocity_to_sim(root_velocity, env_ids=env_ids) + articulation.write_root_link_pose_to_sim(root_pose, env_ids=env_ids) + articulation.write_root_com_velocity_to_sim(root_velocity, env_ids=env_ids) # joint state joint_position = asset_state["joint_position"].clone() joint_velocity = asset_state["joint_velocity"].clone() @@ -463,8 +463,8 @@ def reset_to( if is_relative: root_pose[:, :3] += self.env_origins[env_ids] root_velocity = asset_state["root_velocity"].clone() - rigid_object.write_root_pose_to_sim(root_pose, env_ids=env_ids) - rigid_object.write_root_velocity_to_sim(root_velocity, env_ids=env_ids) + rigid_object.write_root_link_pose_to_sim(root_pose, env_ids=env_ids) + rigid_object.write_root_com_velocity_to_sim(root_velocity, env_ids=env_ids) self.write_data_to_sim() def write_data_to_sim(self): diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/utils/math.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/utils/math.py index 42c6a2bc0d..98e14f5a1c 100644 --- a/source/extensions/omni.isaac.lab/omni/isaac/lab/utils/math.py +++ b/source/extensions/omni.isaac.lab/omni/isaac/lab/utils/math.py @@ -746,7 +746,7 @@ def is_identity_pose(pos: torch.tensor, rot: torch.tensor) -> bool: return torch.allclose(pos, pos_identity) and torch.allclose(rot, rot_identity) -# @torch.jit.script +@torch.jit.script def combine_frame_transforms( t01: torch.Tensor, q01: torch.Tensor, t12: torch.Tensor | None = None, q12: torch.Tensor | None = None ) -> tuple[torch.Tensor, torch.Tensor]: diff --git a/source/extensions/omni.isaac.lab/test/assets/check_external_force.py b/source/extensions/omni.isaac.lab/test/assets/check_external_force.py index ef94da4975..2ba423d239 100644 --- a/source/extensions/omni.isaac.lab/test/assets/check_external_force.py +++ b/source/extensions/omni.isaac.lab/test/assets/check_external_force.py @@ -99,7 +99,8 @@ def main(): root_state = robot.data.default_root_state.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_state_to_sim(root_state) + robot.write_root_link_pose_to_sim(root_state[:, :7]) + robot.write_root_com_velocity_to_sim(root_state[:, 7:]) # reset dof state joint_pos, joint_vel = robot.data.default_joint_pos, robot.data.default_joint_vel robot.write_joint_state_to_sim(joint_pos, joint_vel) diff --git a/source/extensions/omni.isaac.lab/test/assets/check_fixed_base_assets.py b/source/extensions/omni.isaac.lab/test/assets/check_fixed_base_assets.py index b5e43b9016..50b8aeb113 100644 --- a/source/extensions/omni.isaac.lab/test/assets/check_fixed_base_assets.py +++ b/source/extensions/omni.isaac.lab/test/assets/check_fixed_base_assets.py @@ -113,7 +113,8 @@ def run_simulator(sim: sim_utils.SimulationContext, entities: dict[str, Articula root_state = robot.data.default_root_state.clone() root_state[:, :3] += origins[index] root_state[:, :2] += torch.randn_like(root_state[:, :2]) * 0.25 - robot.write_root_state_to_sim(root_state) + robot.write_root_link_pose_to_sim(root_state[:, :7]) + robot.write_root_com_velocity_to_sim(root_state[:, 7:]) # joint state joint_pos, joint_vel = robot.data.default_joint_pos.clone(), robot.data.default_joint_vel.clone() robot.write_joint_state_to_sim(joint_pos, joint_vel) diff --git a/source/extensions/omni.isaac.lab/test/assets/test_articulation.py b/source/extensions/omni.isaac.lab/test/assets/test_articulation.py index d96283be6a..d4b077f9ab 100644 --- a/source/extensions/omni.isaac.lab/test/assets/test_articulation.py +++ b/source/extensions/omni.isaac.lab/test/assets/test_articulation.py @@ -26,6 +26,7 @@ import omni.isaac.core.utils.prims as prim_utils import omni.isaac.lab.sim as sim_utils +import omni.isaac.lab.utils.math as math_utils import omni.isaac.lab.utils.string as string_utils from omni.isaac.lab.actuators import ImplicitActuatorCfg from omni.isaac.lab.assets import Articulation, ArticulationCfg @@ -107,9 +108,9 @@ def generate_articulation( The articulation and environment translations. """ - # Generate translations of 1.5m in x for each articulation + # Generate translations of 2.5 m in x for each articulation translations = torch.zeros(num_articulations, 3, device=device) - translations[:, 0] = torch.arange(num_articulations) * 1.5 + translations[:, 0] = torch.arange(num_articulations) * 2.5 # Create Top-level Xforms, one for each articulation for i in range(num_articulations): @@ -149,8 +150,8 @@ def test_initialization_floating_base_non_root(self): # Check that is fixed base self.assertFalse(articulation.is_fixed_base) # Check buffers that exists and have correct shapes - self.assertEqual(articulation.data.root_pos_w.shape, (num_articulations, 3)) - self.assertEqual(articulation.data.root_quat_w.shape, (num_articulations, 4)) + self.assertEqual(articulation.data.root_link_pos_w.shape, (num_articulations, 3)) + self.assertEqual(articulation.data.root_link_quat_w.shape, (num_articulations, 4)) self.assertEqual(articulation.data.joint_pos.shape, (num_articulations, 21)) # Check some internal physx data for debugging @@ -198,8 +199,8 @@ def test_initialization_floating_base(self): # Check that floating base self.assertFalse(articulation.is_fixed_base) # Check buffers that exists and have correct shapes - self.assertEqual(articulation.data.root_pos_w.shape, (num_articulations, 3)) - self.assertEqual(articulation.data.root_quat_w.shape, (num_articulations, 4)) + self.assertEqual(articulation.data.root_link_pos_w.shape, (num_articulations, 3)) + self.assertEqual(articulation.data.root_link_quat_w.shape, (num_articulations, 4)) self.assertEqual(articulation.data.joint_pos.shape, (num_articulations, 12)) self.assertEqual( articulation.data.default_mass.shape, (num_articulations, articulation.num_bodies) @@ -251,8 +252,8 @@ def test_initialization_fixed_base(self): # Check that fixed base self.assertTrue(articulation.is_fixed_base) # Check buffers that exists and have correct shapes - self.assertEqual(articulation.data.root_pos_w.shape, (num_articulations, 3)) - self.assertEqual(articulation.data.root_quat_w.shape, (num_articulations, 4)) + self.assertEqual(articulation.data.root_link_pos_w.shape, (num_articulations, 3)) + self.assertEqual(articulation.data.root_link_quat_w.shape, (num_articulations, 4)) self.assertEqual(articulation.data.joint_pos.shape, (num_articulations, 9)) self.assertEqual( articulation.data.default_mass.shape, (num_articulations, articulation.num_bodies) @@ -310,8 +311,8 @@ def test_initialization_fixed_base_single_joint(self): # Check that fixed base self.assertTrue(articulation.is_fixed_base) # Check buffers that exists and have correct shapes - self.assertEqual(articulation.data.root_pos_w.shape, (num_articulations, 3)) - self.assertEqual(articulation.data.root_quat_w.shape, (num_articulations, 4)) + self.assertEqual(articulation.data.root_link_pos_w.shape, (num_articulations, 3)) + self.assertEqual(articulation.data.root_link_quat_w.shape, (num_articulations, 4)) self.assertEqual(articulation.data.joint_pos.shape, (num_articulations, 1)) self.assertEqual( articulation.data.default_mass.shape, (num_articulations, articulation.num_bodies) @@ -369,8 +370,8 @@ def test_initialization_hand_with_tendons(self): # Check that fixed base self.assertTrue(articulation.is_fixed_base) # Check buffers that exists and have correct shapes - self.assertTrue(articulation.data.root_pos_w.shape == (num_articulations, 3)) - self.assertTrue(articulation.data.root_quat_w.shape == (num_articulations, 4)) + self.assertTrue(articulation.data.root_link_pos_w.shape == (num_articulations, 3)) + self.assertTrue(articulation.data.root_link_quat_w.shape == (num_articulations, 4)) self.assertTrue(articulation.data.joint_pos.shape == (num_articulations, 24)) self.assertEqual( articulation.data.default_mass.shape, (num_articulations, articulation.num_bodies) @@ -419,8 +420,8 @@ def test_initialization_floating_base_made_fixed_base(self): # Check that is fixed base self.assertTrue(articulation.is_fixed_base) # Check buffers that exists and have correct shapes - self.assertEqual(articulation.data.root_pos_w.shape, (num_articulations, 3)) - self.assertEqual(articulation.data.root_quat_w.shape, (num_articulations, 4)) + self.assertEqual(articulation.data.root_link_pos_w.shape, (num_articulations, 3)) + self.assertEqual(articulation.data.root_link_quat_w.shape, (num_articulations, 4)) self.assertEqual(articulation.data.joint_pos.shape, (num_articulations, 12)) # Check some internal physx data for debugging @@ -474,8 +475,8 @@ def test_initialization_fixed_base_made_floating_base(self): # Check that is floating base self.assertFalse(articulation.is_fixed_base) # Check buffers that exists and have correct shapes - self.assertEqual(articulation.data.root_pos_w.shape, (num_articulations, 3)) - self.assertEqual(articulation.data.root_quat_w.shape, (num_articulations, 4)) + self.assertEqual(articulation.data.root_link_pos_w.shape, (num_articulations, 3)) + self.assertEqual(articulation.data.root_link_quat_w.shape, (num_articulations, 4)) self.assertEqual(articulation.data.joint_pos.shape, (num_articulations, 9)) # Check some internal physx data for debugging @@ -632,7 +633,8 @@ def test_external_force_on_single_body(self): # reset root state root_state = articulation.data.default_root_state.clone() - articulation.write_root_state_to_sim(root_state) + articulation.write_root_link_pose_to_sim(root_state[:, :7]) + articulation.write_root_com_velocity_to_sim(root_state[:, 7:]) # reset dof state joint_pos, joint_vel = ( articulation.data.default_joint_pos, @@ -656,7 +658,7 @@ def test_external_force_on_single_body(self): articulation.update(sim.cfg.dt) # check condition that the articulations have fallen down for i in range(num_articulations): - self.assertLess(articulation.data.root_pos_w[i, 2].item(), 0.2) + self.assertLess(articulation.data.root_link_pos_w[i, 2].item(), 0.2) def test_external_force_on_multiple_bodies(self): """Test application of external force on the legs of the articulation.""" @@ -679,7 +681,12 @@ def test_external_force_on_multiple_bodies(self): # Now we are ready! for _ in range(5): # reset root state - articulation.write_root_state_to_sim(articulation.data.default_root_state.clone()) + articulation.write_root_link_pose_to_sim( + articulation.data.default_root_state.clone()[:, :7] + ) + articulation.write_root_com_velocity_to_sim( + articulation.data.default_root_state.clone()[:, 7:] + ) # reset dof state joint_pos, joint_vel = ( articulation.data.default_joint_pos, @@ -704,7 +711,7 @@ def test_external_force_on_multiple_bodies(self): # check condition for i in range(num_articulations): # since there is a moment applied on the articulation, the articulation should rotate - self.assertTrue(articulation.data.root_ang_vel_w[i, 2].item() > 0.1) + self.assertTrue(articulation.data.root_com_ang_vel_w[i, 2].item() > 0.1) def test_loading_gains_from_usd(self): """Test that gains are loaded from USD file if actuator model has them as None.""" @@ -922,6 +929,206 @@ def test_apply_joint_command(self): # are not properly tuned assert not torch.allclose(articulation.data.joint_pos, joint_pos) + def test_body_root_link_state(self): + """Test for the root_link_state_w property""" + for num_articulations in (1, 2): + # for num_articulations in ( 2,): + for device in ("cuda:0", "cpu"): + # for device in ("cuda:0",): + for with_offset in [True, False]: + # for with_offset in [True,]: + with self.subTest(num_articulations=num_articulations, device=device, with_offset=with_offset): + with build_simulation_context( + device=device, add_ground_plane=False, auto_add_lighting=True + ) as sim: + articulation_cfg = generate_articulation_cfg(articulation_type="single_joint") + articulation, env_pos = generate_articulation(articulation_cfg, num_articulations, device) + env_idx = torch.tensor([x for x in range(num_articulations)]) + # Check that boundedness of articulation is correct + self.assertEqual(ctypes.c_long.from_address(id(articulation)).value, 1) + # Play sim + sim.reset() + # Check if articulation is initialized + self.assertTrue(articulation.is_initialized) + # Check that fixed base + self.assertTrue(articulation.is_fixed_base) + + # change center of mass offset from link frame + if with_offset: + offset = [0.5, 0.0, 0.0] + else: + offset = [0.0, 0.0, 0.0] + + # create com offsets + num_bodies = articulation.num_bodies + com = articulation.root_physx_view.get_coms() + link_offset = [1.0, 0.0, 0.0] # the offset from CenterPivot to Arm frames + new_com = torch.tensor(offset, device=device).repeat(num_articulations, 1, 1) + com[:, 1, :3] = new_com.squeeze(-2) + articulation.root_physx_view.set_coms(com, env_idx) + + # check they are set + torch.testing.assert_close(articulation.root_physx_view.get_coms(), com) + + for i in range(50): + # perform step + sim.step() + # update buffers + articulation.update(sim.cfg.dt) + + # get state properties + root_state_w = articulation.data.root_state_w + root_link_state_w = articulation.data.root_link_state_w + root_com_state_w = articulation.data.root_com_state_w + body_state_w = articulation.data.body_state_w + body_link_state_w = articulation.data.body_link_state_w + body_com_state_w = articulation.data.body_com_state_w + + if with_offset: + # get joint state + joint_pos = articulation.data.joint_pos.unsqueeze(-1) + joint_vel = articulation.data.joint_vel.unsqueeze(-1) + + # LINK state + # pose + torch.testing.assert_close(root_state_w[..., :7], root_link_state_w[..., :7]) + torch.testing.assert_close(body_state_w[..., :7], body_link_state_w[..., :7]) + + # lin_vel arm + lin_vel_gt = torch.zeros(num_articulations, num_bodies, 3, device=device) + vx = -(link_offset[0]) * joint_vel * torch.sin(joint_pos) + vy = torch.zeros(num_articulations, 1, 1, device=device) + vz = (link_offset[0]) * joint_vel * torch.cos(joint_pos) + lin_vel_gt[:, 1, :] = torch.cat([vx, vy, vz], dim=-1).squeeze(-2) + + # linear velocity of root link should be zero + torch.testing.assert_close( + lin_vel_gt[:, 0, :], root_link_state_w[..., 7:10], atol=1e-3, rtol=1e-1 + ) + # linear velocity of pendulum link should be + torch.testing.assert_close( + lin_vel_gt, body_link_state_w[..., 7:10], atol=1e-3, rtol=1e-1 + ) + + # ang_vel + torch.testing.assert_close(root_state_w[..., 10:], root_link_state_w[..., 10:]) + torch.testing.assert_close(body_state_w[..., 10:], body_link_state_w[..., 10:]) + + # COM state + # position and orientation shouldn't match for the _state_com_w but everything else will + pos_gt = torch.zeros(num_articulations, num_bodies, 3, device=device) + px = (link_offset[0] + offset[0]) * torch.cos(joint_pos) + py = torch.zeros(num_articulations, 1, 1, device=device) + pz = (link_offset[0] + offset[0]) * torch.sin(joint_pos) + pos_gt[:, 1, :] = torch.cat([px, py, pz], dim=-1).squeeze(-2) + pos_gt += env_pos.unsqueeze(-2).repeat(1, num_bodies, 1) + torch.testing.assert_close( + pos_gt[:, 0, :], root_com_state_w[..., :3], atol=1e-3, rtol=1e-1 + ) + 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_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]) + + # linear vel, and angular vel + torch.testing.assert_close(root_state_w[..., 7:], root_com_state_w[..., 7:]) + torch.testing.assert_close(body_state_w[..., 7:], body_com_state_w[..., 7:]) + else: + # single joint center of masses are at link frames so they will be the same + torch.testing.assert_close(root_state_w, root_link_state_w) + torch.testing.assert_close(root_state_w, root_com_state_w) + torch.testing.assert_close(body_state_w, body_link_state_w) + torch.testing.assert_close(body_state_w, body_com_state_w) + + def test_write_root_state(self): + """Test the setters for root_state using both the link frame and center of mass as reference frame.""" + for num_articulations in (1, 2): + for device in ("cuda:0", "cpu"): + for with_offset in [True, False]: + for state_location in ("com", "link"): + with self.subTest( + num_articulations=num_articulations, + device=device, + with_offset=with_offset, + state_location=state_location, + ): + with build_simulation_context( + device=device, add_ground_plane=False, auto_add_lighting=True, gravity_enabled=False + ) as sim: + articulation_cfg = generate_articulation_cfg(articulation_type="anymal") + articulation, env_pos = generate_articulation( + articulation_cfg, num_articulations, device + ) + env_idx = torch.tensor([x for x in range(num_articulations)]) + + # Play sim + sim.reset() + + # change center of mass offset from link frame + if with_offset: + offset = torch.tensor([1.0, 0.0, 0.0]).repeat(num_articulations, 1, 1) + else: + offset = torch.tensor([0.0, 0.0, 0.0]).repeat(num_articulations, 1, 1) + + # create com offsets + com = articulation.root_physx_view.get_coms() + new_com = offset + com[:, 0, :3] = new_com.squeeze(-2) + articulation.root_physx_view.set_coms(com, env_idx) + + # check they are set + torch.testing.assert_close(articulation.root_physx_view.get_coms(), com) + + rand_state = torch.zeros_like(articulation.data.root_link_state_w) + rand_state[..., :7] = articulation.data.default_root_state[..., :7] + rand_state[..., :3] += env_pos + # make quaternion a unit vector + rand_state[..., 3:7] = torch.nn.functional.normalize(rand_state[..., 3:7], dim=-1) + + for i in range(10): + + # perform step + sim.step() + # update buffers + articulation.update(sim.cfg.dt) + + if state_location == "com": + articulation.write_root_com_state_to_sim(rand_state) + elif state_location == "link": + articulation.write_root_link_state_to_sim(rand_state) + + if state_location == "com": + torch.testing.assert_close(rand_state, articulation.data.root_com_state_w) + elif state_location == "link": + torch.testing.assert_close(rand_state, articulation.data.root_link_state_w) + + def test_transform_inverses(self): + """Test if math utilities are true inverses of each other.""" + + pose01 = torch.rand(1, 7) + pose01[:, 3:7] = torch.nn.functional.normalize(pose01[..., 3:7], dim=-1) + + pose12 = torch.rand(1, 7) + pose12[:, 3:7] = torch.nn.functional.normalize(pose12[..., 3:7], dim=-1) + + pos02, quat02 = math_utils.combine_frame_transforms( + pose01[..., :3], pose01[..., 3:7], pose12[:, :3], pose12[:, 3:7] + ) + + pos01, quat01 = math_utils.combine_frame_transforms( + pos02, + quat02, + math_utils.quat_rotate(math_utils.quat_inv(pose12[:, 3:7]), -pose12[:, :3]), + math_utils.quat_inv(pose12[:, 3:7]), + ) + print("") + print(pose01) + print(pos01, quat01) + torch.testing.assert_close(pose01, torch.cat((pos01, quat01), dim=-1)) + if __name__ == "__main__": run_tests() diff --git a/source/extensions/omni.isaac.lab/test/assets/test_rigid_object.py b/source/extensions/omni.isaac.lab/test/assets/test_rigid_object.py index 2724f13e05..06d3087d7a 100644 --- a/source/extensions/omni.isaac.lab/test/assets/test_rigid_object.py +++ b/source/extensions/omni.isaac.lab/test/assets/test_rigid_object.py @@ -33,7 +33,7 @@ from omni.isaac.lab.sim import build_simulation_context from omni.isaac.lab.sim.spawners import materials from omni.isaac.lab.utils.assets import ISAAC_NUCLEUS_DIR, ISAACLAB_NUCLEUS_DIR -from omni.isaac.lab.utils.math import default_orientation, random_orientation +from omni.isaac.lab.utils.math import default_orientation, quat_mul, quat_rotate_inverse, random_orientation def generate_cubes_scene( @@ -119,8 +119,8 @@ def test_initialization(self): self.assertEqual(len(cube_object.body_names), 1) # Check buffers that exists and have correct shapes - self.assertEqual(cube_object.data.root_pos_w.shape, (num_cubes, 3)) - self.assertEqual(cube_object.data.root_quat_w.shape, (num_cubes, 4)) + self.assertEqual(cube_object.data.root_link_pos_w.shape, (num_cubes, 3)) + self.assertEqual(cube_object.data.root_link_quat_w.shape, (num_cubes, 4)) self.assertEqual(cube_object.data.default_mass.shape, (num_cubes, 1)) self.assertEqual(cube_object.data.default_inertia.shape, (num_cubes, 9)) @@ -153,8 +153,8 @@ def test_initialization_with_kinematic_enabled(self): self.assertEqual(len(cube_object.body_names), 1) # Check buffers that exists and have correct shapes - self.assertEqual(cube_object.data.root_pos_w.shape, (num_cubes, 3)) - self.assertEqual(cube_object.data.root_quat_w.shape, (num_cubes, 4)) + self.assertEqual(cube_object.data.root_link_pos_w.shape, (num_cubes, 3)) + self.assertEqual(cube_object.data.root_link_quat_w.shape, (num_cubes, 4)) # Simulate physics for _ in range(2): @@ -237,7 +237,8 @@ def test_external_force_on_single_body(self): # need to shift the position of the cubes otherwise they will be on top of each other root_state[:, :3] = origins - cube_object.write_root_state_to_sim(root_state) + cube_object.write_root_link_pose_to_sim(root_state[:, :7]) + cube_object.write_root_com_velocity_to_sim(root_state[:, 7:]) # reset object cube_object.reset() @@ -259,10 +260,10 @@ def test_external_force_on_single_body(self): # First object should still be at the same Z position (1.0) torch.testing.assert_close( - cube_object.data.root_pos_w[0::2, 2], torch.ones(num_cubes // 2, device=sim.device) + cube_object.data.root_link_pos_w[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 - self.assertTrue(torch.all(cube_object.data.root_pos_w[1::2, 2] < 1.0)) + self.assertTrue(torch.all(cube_object.data.root_link_pos_w[1::2, 2] < 1.0)) def test_set_rigid_object_state(self): """Test setting the state of the rigid object. @@ -288,10 +289,14 @@ def test_set_rigid_object_state(self): # 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(cube_object.data.root_pos_w, device=sim.device), + "root_pos_w": torch.zeros_like(cube_object.data.root_link_pos_w, device=sim.device), "root_quat_w": default_orientation(num=num_cubes, device=sim.device), - "root_lin_vel_w": torch.zeros_like(cube_object.data.root_lin_vel_w, device=sim.device), - "root_ang_vel_w": torch.zeros_like(cube_object.data.root_ang_vel_w, device=sim.device), + "root_lin_vel_w": torch.zeros_like( + cube_object.data.root_com_lin_vel_w, device=sim.device + ), + "root_ang_vel_w": torch.zeros_like( + cube_object.data.root_com_ang_vel_w, device=sim.device + ), } # Now we are ready! @@ -319,7 +324,8 @@ def test_set_rigid_object_state(self): dim=-1, ) # reset root state - cube_object.write_root_state_to_sim(root_state=root_state) + cube_object.write_root_link_pose_to_sim(root_state[:, :7]) + cube_object.write_root_com_velocity_to_sim(root_state[:, 7:]) sim.step() @@ -355,7 +361,8 @@ def test_reset_rigid_object(self): # Random orientation root_state[:, 3:7] = random_orientation(num=num_cubes, device=sim.device) - cube_object.write_root_state_to_sim(root_state) + cube_object.write_root_link_pose_to_sim(root_state[:, :7]) + cube_object.write_root_com_velocity_to_sim(root_state[:, 7:]) if i % 2 == 0: # reset object @@ -440,7 +447,7 @@ def test_rigid_body_no_friction(self): initial_velocity = torch.zeros((num_cubes, 6), device=sim.cfg.device) initial_velocity[:, 0] = 0.1 - cube_object.write_root_velocity_to_sim(initial_velocity) + cube_object.write_root_com_velocity_to_sim(initial_velocity) # Simulate physics for _ in range(5): @@ -505,7 +512,7 @@ def test_rigid_body_with_static_friction(self): for _ in range(100): sim.step() cube_object.update(sim.cfg.dt) - cube_object.write_root_velocity_to_sim(torch.zeros((num_cubes, 6), device=sim.device)) + cube_object.write_root_com_velocity_to_sim(torch.zeros((num_cubes, 6), device=sim.device)) cube_mass = cube_object.root_physx_view.get_masses() gravity_magnitude = abs(sim.cfg.gravity[2]) # 2 cases: force applied is below and above mu @@ -514,7 +521,9 @@ def test_rigid_body_with_static_friction(self): for force in "below_mu", "above_mu": with self.subTest(force=force): # set initial velocity to zero - cube_object.write_root_velocity_to_sim(torch.zeros((num_cubes, 6), device=sim.device)) + cube_object.write_root_com_velocity_to_sim( + torch.zeros((num_cubes, 6), device=sim.device) + ) external_wrench_b = torch.zeros((num_cubes, 1, 6), device=sim.device) if force == "below_mu": @@ -532,7 +541,7 @@ def test_rigid_body_with_static_friction(self): ) # Get root state - initial_root_pos = cube_object.data.root_pos_w.clone() + initial_root_pos = cube_object.data.root_link_pos_w.clone() # Simulate physics for _ in range(200): # apply the wrench @@ -543,7 +552,7 @@ def test_rigid_body_with_static_friction(self): if force == "below_mu": # Assert that the block has not moved torch.testing.assert_close( - cube_object.data.root_pos_w, initial_root_pos, rtol=1e-3, atol=1e-3 + cube_object.data.root_link_pos_w, initial_root_pos, rtol=1e-3, atol=1e-3 ) if force == "above_mu": self.assertTrue( @@ -595,7 +604,8 @@ def test_rigid_body_with_restitution(self): root_state[:, 2] = 1.0 # Set an initial drop height root_state[:, 9] = -1.0 # Set an initial downward velocity - cube_object.write_root_state_to_sim(root_state=root_state) + cube_object.write_root_link_pose_to_sim(root_state[:, :7]) + cube_object.write_root_com_velocity_to_sim(root_state[:, 7:]) static_friction = torch.zeros(num_cubes, 1) dynamic_friction = torch.zeros(num_cubes, 1) @@ -606,14 +616,14 @@ def test_rigid_body_with_restitution(self): # Add restitution to cube cube_object.root_physx_view.set_material_properties(cube_object_materials, indices) - curr_z_velocity = cube_object.data.root_lin_vel_w[:, 2].clone() + curr_z_velocity = cube_object.data.root_com_lin_vel_w[:, 2].clone() for _ in range(100): sim.step() # update object cube_object.update(sim.cfg.dt) - curr_z_velocity = cube_object.data.root_lin_vel_w[:, 2].clone() + curr_z_velocity = cube_object.data.root_com_lin_vel_w[:, 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 @@ -709,6 +719,182 @@ def test_gravity_vec_w(self): # Check the body accelerations are correct torch.testing.assert_close(cube_object.data.body_acc_w, gravity) + def test_body_root_state_properties(self): + """Test the root_com_state_w, root_link_state_w, body_com_state_w, and body_link_state_w properties.""" + for num_cubes in (1, 2): + for device in ("cuda:0", "cpu"): + for with_offset in [True, False]: + with self.subTest(num_cubes=num_cubes, device=device, with_offset=with_offset): + with build_simulation_context( + device=device, gravity_enabled=False, auto_add_lighting=True + ) as sim: + # Create a scene with random cubes + cube_object, env_pos = generate_cubes_scene(num_cubes=num_cubes, height=0.0, device=device) + env_idx = torch.tensor([x for x in range(num_cubes)]) + + # Play sim + sim.reset() + + # Check if cube_object is initialized + self.assertTrue(cube_object.is_initialized) + + # change center of mass offset from link frame + if with_offset: + offset = torch.tensor([0.1, 0.0, 0.0], device=device).repeat(num_cubes, 1) + else: + offset = torch.tensor([0.0, 0.0, 0.0], device=device).repeat(num_cubes, 1) + + com = cube_object.root_physx_view.get_coms() + com[..., :3] = offset.to("cpu") + cube_object.root_physx_view.set_coms(com, env_idx) + + # check ceter of mass has been set + torch.testing.assert_close(cube_object.root_physx_view.get_coms(), com) + + # random z spin velocity + spin_twist = torch.zeros(6, device=device) + spin_twist[5] = torch.randn(1, device=device) + + # Simulate physics + for _ in range(100): + # spin the object around Z axis (com) + cube_object.write_root_com_velocity_to_sim(spin_twist.repeat(num_cubes, 1)) + # perform rendering + sim.step() + # update object + cube_object.update(sim.cfg.dt) + + # get state properties + root_state_w = cube_object.data.root_state_w + root_link_state_w = cube_object.data.root_link_state_w + root_com_state_w = cube_object.data.root_com_state_w + body_state_w = cube_object.data.body_state_w + body_link_state_w = cube_object.data.body_link_state_w + body_com_state_w = cube_object.data.body_com_state_w + + # if offset is [0,0,0] all root_state_%_w will match and all body_%_w will match + if not with_offset: + torch.testing.assert_close(root_state_w, root_com_state_w) + torch.testing.assert_close(root_state_w, root_link_state_w) + torch.testing.assert_close(body_state_w, body_com_state_w) + torch.testing.assert_close(body_state_w, body_link_state_w) + else: + # cubes are spinning around center of mass + # position will not match + # center of mass position will be constant (i.e. spining 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 + root_link_state_pos_rel_com = quat_rotate_inverse( + root_link_state_w[..., 3:7], + root_link_state_w[..., :3] - root_com_state_w[..., :3], + ) + torch.testing.assert_close(-offset, root_link_state_pos_rel_com) + body_link_state_pos_rel_com = quat_rotate_inverse( + body_link_state_w[..., 3:7], + body_link_state_w[..., :3] - body_com_state_w[..., :3], + ) + 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_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]) + + # orientation of link will match root state will always match + torch.testing.assert_close(root_state_w[..., 3:7], root_link_state_w[..., 3:7]) + 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) + 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 + lin_vel_rel_root_gt = quat_rotate_inverse( + root_link_state_w[..., 3:7], root_link_state_w[..., 7:10] + ) + lin_vel_rel_body_gt = quat_rotate_inverse( + body_link_state_w[..., 3:7], body_link_state_w[..., 7:10] + ) + lin_vel_rel_gt = torch.linalg.cross( + spin_twist.repeat(num_cubes, 1)[..., 3:], -offset + ) + torch.testing.assert_close( + lin_vel_rel_gt, lin_vel_rel_root_gt, atol=1e-4, rtol=1e-4 + ) + torch.testing.assert_close( + lin_vel_rel_gt, lin_vel_rel_body_gt.squeeze(-2), atol=1e-4, rtol=1e-4 + ) + + # ang_vel will always match + torch.testing.assert_close(root_state_w[..., 10:], root_com_state_w[..., 10:]) + torch.testing.assert_close(root_state_w[..., 10:], root_link_state_w[..., 10:]) + torch.testing.assert_close(body_state_w[..., 10:], body_com_state_w[..., 10:]) + torch.testing.assert_close(body_state_w[..., 10:], body_link_state_w[..., 10:]) + + def test_write_root_state(self): + """Test the setters for root_state using both the link frame and center of mass as reference frame.""" + for num_cubes in (1, 2): + for device in ("cuda:0", "cpu"): + for with_offset in [True, False]: + for state_location in ("com", "link"): + with self.subTest(num_cubes=num_cubes, device=device, with_offset=with_offset): + with build_simulation_context( + device=device, gravity_enabled=False, auto_add_lighting=True + ) as sim: + # Create a scene with random cubes + cube_object, env_pos = generate_cubes_scene( + num_cubes=num_cubes, height=0.0, device=device + ) + env_idx = torch.tensor([x for x in range(num_cubes)]) + + # Play sim + sim.reset() + + # Check if cube_object is initialized + self.assertTrue(cube_object.is_initialized) + + # change center of mass offset from link frame + if with_offset: + offset = torch.tensor([0.1, 0.0, 0.0], device=device).repeat(num_cubes, 1) + else: + offset = torch.tensor([0.0, 0.0, 0.0], device=device).repeat(num_cubes, 1) + + com = cube_object.root_physx_view.get_coms() + com[..., :3] = offset.to("cpu") + cube_object.root_physx_view.set_coms(com, env_idx) + + # check ceter 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_link_state_w) + rand_state[..., :7] = cube_object.data.default_root_state[..., :7] + rand_state[..., :3] += env_pos + # make quaternion a unit vector + rand_state[..., 3:7] = torch.nn.functional.normalize(rand_state[..., 3:7], dim=-1) + + for i in range(10): + + # perform step + sim.step() + # update buffers + cube_object.update(sim.cfg.dt) + + if state_location == "com": + cube_object.write_root_com_state_to_sim(rand_state) + elif state_location == "link": + cube_object.write_root_link_state_to_sim(rand_state) + + if state_location == "com": + torch.testing.assert_close(rand_state, cube_object.data.root_com_state_w) + elif state_location == "link": + torch.testing.assert_close(rand_state, cube_object.data.root_link_state_w) + if __name__ == "__main__": run_tests() diff --git a/source/extensions/omni.isaac.lab/test/assets/test_rigid_object_collection.py b/source/extensions/omni.isaac.lab/test/assets/test_rigid_object_collection.py index 5674359ab0..44b5a570fe 100644 --- a/source/extensions/omni.isaac.lab/test/assets/test_rigid_object_collection.py +++ b/source/extensions/omni.isaac.lab/test/assets/test_rigid_object_collection.py @@ -31,7 +31,7 @@ from omni.isaac.lab.assets import RigidObjectCfg, RigidObjectCollection, RigidObjectCollectionCfg from omni.isaac.lab.sim import build_simulation_context from omni.isaac.lab.utils.assets import ISAAC_NUCLEUS_DIR -from omni.isaac.lab.utils.math import default_orientation, random_orientation +from omni.isaac.lab.utils.math import default_orientation, quat_mul, quat_rotate_inverse, random_orientation def generate_cubes_scene( @@ -120,8 +120,8 @@ def test_initialization(self): self.assertEqual(len(object_collection.object_names), num_cubes) # Check buffers that exists and have correct shapes - self.assertEqual(object_collection.data.object_pos_w.shape, (num_envs, num_cubes, 3)) - self.assertEqual(object_collection.data.object_quat_w.shape, (num_envs, num_cubes, 4)) + self.assertEqual(object_collection.data.object_link_pos_w.shape, (num_envs, num_cubes, 3)) + self.assertEqual(object_collection.data.object_link_quat_w.shape, (num_envs, num_cubes, 4)) self.assertEqual(object_collection.data.default_mass.shape, (num_envs, num_cubes, 1)) self.assertEqual(object_collection.data.default_inertia.shape, (num_envs, num_cubes, 9)) @@ -190,8 +190,8 @@ def test_initialization_with_kinematic_enabled(self): self.assertEqual(len(object_collection.object_names), num_cubes) # Check buffers that exists and have correct shapes - self.assertEqual(object_collection.data.object_pos_w.shape, (num_envs, num_cubes, 3)) - self.assertEqual(object_collection.data.object_quat_w.shape, (num_envs, num_cubes, 4)) + self.assertEqual(object_collection.data.object_link_pos_w.shape, (num_envs, num_cubes, 3)) + self.assertEqual(object_collection.data.object_link_quat_w.shape, (num_envs, num_cubes, 4)) # Simulate physics for _ in range(2): @@ -202,7 +202,9 @@ def test_initialization_with_kinematic_enabled(self): # check that the object is kinematic default_object_state = object_collection.data.default_object_state.clone() default_object_state[..., :3] += origins.unsqueeze(1) - torch.testing.assert_close(object_collection.data.object_state_w, default_object_state) + torch.testing.assert_close( + object_collection.data.object_link_state_w, default_object_state + ) def test_initialization_with_no_rigid_body(self): """Test that initialization fails when no rigid body is found at the provided prim path.""" @@ -283,11 +285,11 @@ def test_external_force_on_single_body(self): # First object should still be at the same Z position (1.0) torch.testing.assert_close( - object_collection.data.object_pos_w[:, 0::2, 2], + object_collection.data.object_link_pos_w[:, 0::2, 2], torch.ones_like(object_collection.data.object_pos_w[:, 0::2, 2]), ) # Second object should have fallen, so it's Z height should be less than initial height of 1.0 - self.assertTrue(torch.all(object_collection.data.object_pos_w[:, 1::2, 2] < 1.0)) + self.assertTrue(torch.all(object_collection.data.object_link_pos_w[:, 1::2, 2] < 1.0)) def test_set_object_state(self): """Test setting the state of the object. @@ -375,6 +377,193 @@ def test_set_object_state(self): object_collection.update(sim.cfg.dt) + def test_object_state_properties(self): + """Test the object_com_state_w and object_link_state_w properties.""" + for num_envs in (1, 4): + for num_cubes in (1, 2): + for device in ("cuda:0", "cpu"): + for with_offset in [True, False]: + with self.subTest( + num_envs=num_envs, num_cubes=num_cubes, device=device, with_offset=with_offset + ): + with build_simulation_context( + device=device, gravity_enabled=False, auto_add_lighting=True + ) as sim: + # Create a scene with random cubes + cube_object, env_pos = generate_cubes_scene( + num_envs=num_envs, num_cubes=num_cubes, height=0.0, device=device + ) + view_ids = torch.tensor([x for x in range(num_cubes * num_envs)]) + + # Play sim + sim.reset() + + # Check if cube_object is initialized + self.assertTrue(cube_object.is_initialized) + + # change center of mass offset from link frame + if with_offset: + offset = torch.tensor([0.1, 0.0, 0.0], device=device).repeat(num_envs, num_cubes, 1) + else: + offset = torch.tensor([0.0, 0.0, 0.0], device=device).repeat(num_envs, num_cubes, 1) + + com = cube_object.reshape_view_to_data(cube_object.root_physx_view.get_coms()) + com[..., :3] = offset.to("cpu") + cube_object.root_physx_view.set_coms( + cube_object.reshape_data_to_view(com.clone()), view_ids + ) + + # check center of mass has been set + torch.testing.assert_close( + cube_object.reshape_view_to_data(cube_object.root_physx_view.get_coms()), com + ) + + # random z spin velocity + spin_twist = torch.zeros(6, device=device) + spin_twist[5] = torch.randn(1, device=device) + + # initial spawn point + init_com = cube_object.data.object_com_state_w[..., :3] + + # Simulate physics + 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) + ) + # perform rendering + sim.step() + # update object + cube_object.update(sim.cfg.dt) + + # get state properties + object_state_w = cube_object.data.object_state_w + object_link_state_w = cube_object.data.object_link_state_w + object_com_state_w = cube_object.data.object_com_state_w + + # if offset is [0,0,0] all object_state_%_w will match and all body_%_w will match + if not with_offset: + torch.testing.assert_close(object_state_w, object_com_state_w) + torch.testing.assert_close(object_state_w, object_link_state_w) + else: + # cubes are spinning around center of mass + # position will not match + # center of mass position will be constant (i.e. spinning around com) + torch.testing.assert_close(init_com, object_com_state_w[..., :3]) + + # link position will be moving but should stay constant away from center of mass + object_link_state_pos_rel_com = quat_rotate_inverse( + object_link_state_w[..., 3:7], + object_link_state_w[..., :3] - object_com_state_w[..., :3], + ) + + 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_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]) + + # orientation of link will match object state will always match + torch.testing.assert_close( + object_state_w[..., 3:7], object_link_state_w[..., 3:7] + ) + + # lin_vel will not match + # center of mass vel will be constant (i.e. spining around com) + torch.testing.assert_close( + torch.zeros_like(object_com_state_w[..., 7:10]), + object_com_state_w[..., 7:10], + ) + + # link frame will be moving, and should be equal to input angular velocity cross offset + lin_vel_rel_object_gt = quat_rotate_inverse( + object_link_state_w[..., 3:7], object_link_state_w[..., 7:10] + ) + lin_vel_rel_gt = torch.linalg.cross( + spin_twist.repeat(num_envs, num_cubes, 1)[..., 3:], -offset + ) + torch.testing.assert_close( + lin_vel_rel_gt, lin_vel_rel_object_gt, atol=1e-4, rtol=1e-3 + ) + + # ang_vel will always match + torch.testing.assert_close( + object_state_w[..., 10:], object_com_state_w[..., 10:] + ) + torch.testing.assert_close( + object_state_w[..., 10:], object_link_state_w[..., 10:] + ) + + def test_write_object_state(self): + """Test the setters for object_state using both the link frame and center of mass as reference frame.""" + for num_envs in (1, 3): + for num_cubes in (1, 2): + for device in ("cuda:0", "cpu"): + for with_offset in [True, False]: + for state_location in ("com", "link"): + with self.subTest( + num_envs=num_envs, num_cubes=num_cubes, device=device, with_offset=with_offset + ): + with build_simulation_context( + device=device, gravity_enabled=False, auto_add_lighting=True + ) as sim: + # Create a scene with random cubes + cube_object, env_pos = generate_cubes_scene( + num_envs=num_envs, num_cubes=num_cubes, height=0.0, device=device + ) + view_ids = torch.tensor([x for x in range(num_cubes * num_cubes)]) + + # Play sim + sim.reset() + + # Check if cube_object is initialized + self.assertTrue(cube_object.is_initialized) + + # change center of mass offset from link frame + if with_offset: + offset = torch.tensor([0.1, 0.0, 0.0], device=device).repeat( + num_envs, num_cubes, 1 + ) + else: + offset = torch.tensor([0.0, 0.0, 0.0], device=device).repeat( + num_envs, num_cubes, 1 + ) + + com = cube_object.reshape_view_to_data(cube_object.root_physx_view.get_coms()) + com[..., :3] = offset.to("cpu") + cube_object.root_physx_view.set_coms( + cube_object.reshape_data_to_view(com.clone()), view_ids + ) + + # check center of mass has been set + torch.testing.assert_close( + cube_object.reshape_view_to_data(cube_object.root_physx_view.get_coms()), com + ) + + rand_state = torch.zeros_like(cube_object.data.object_link_state_w) + rand_state[..., :7] = cube_object.data.default_object_state[..., :7] + rand_state[..., :3] += cube_object.data.object_link_pos_w + # make quaternion a unit vector + rand_state[..., 3:7] = torch.nn.functional.normalize(rand_state[..., 3:7], dim=-1) + + for i in range(10): + + # perform step + sim.step() + # update buffers + cube_object.update(sim.cfg.dt) + + if state_location == "com": + cube_object.write_object_com_state_to_sim(rand_state) + elif state_location == "link": + cube_object.write_object_link_state_to_sim(rand_state) + + if state_location == "com": + torch.testing.assert_close(rand_state, cube_object.data.object_com_state_w) + elif state_location == "link": + torch.testing.assert_close(rand_state, cube_object.data.object_link_state_w) + def test_reset_object_collection(self): """Test resetting the state of the rigid object.""" for num_envs in (1, 3): diff --git a/source/extensions/omni.isaac.lab/test/controllers/test_differential_ik.py b/source/extensions/omni.isaac.lab/test/controllers/test_differential_ik.py index 17a069d5b9..afbbb3214c 100644 --- a/source/extensions/omni.isaac.lab/test/controllers/test_differential_ik.py +++ b/source/extensions/omni.isaac.lab/test/controllers/test_differential_ik.py @@ -147,8 +147,8 @@ def _run_ik_controller( ee_pose_b_des = torch.zeros(self.num_envs, diff_ik_controller.action_dim, device=self.sim.device) ee_pose_b_des[:] = self.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_link_state_w[:, ee_frame_idx, 0:7] + root_pose_w = robot.data.root_link_state_w[:, 0:7] 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] ) @@ -177,9 +177,10 @@ def _run_ik_controller( robot.set_joint_position_target(joint_pos) robot.write_data_to_sim() # randomize root state yaw, ik should work regardless base rotation - root_state = robot.data.root_state_w.clone() + root_state = robot.data.root_link_state_w.clone() root_state[:, 3:7] = random_yaw_orientation(self.num_envs, self.sim.device) - robot.write_root_state_to_sim(root_state) + robot.write_root_link_pose_to_sim(root_state[:, :7]) + robot.write_root_com_velocity_to_sim(root_state[:, 7:]) robot.reset() # reset actions ee_pose_b_des[:] = self.ee_pose_b_des_set[current_goal_idx] @@ -194,8 +195,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_link_state_w[:, ee_frame_idx, 0:7] + root_pose_w = robot.data.root_link_state_w[:, 0:7] 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, :]) diff --git a/source/extensions/omni.isaac.lab/test/controllers/test_operational_space.py b/source/extensions/omni.isaac.lab/test/controllers/test_operational_space.py index 151e42c155..bfde34681e 100644 --- a/source/extensions/omni.isaac.lab/test/controllers/test_operational_space.py +++ b/source/extensions/omni.isaac.lab/test/controllers/test_operational_space.py @@ -723,24 +723,26 @@ def _update_states( gravity = robot.root_physx_view.get_generalized_gravity_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(robot.data.root_state_w[:, 3:7])) + root_rot_matrix = matrix_from_quat(quat_inv(robot.data.root_link_quat_w)) 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 = 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_link_state_w[:, 0:7] + ee_pose_w = robot.data.body_link_state_w[:, ee_frame_idx, 0:7] 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 = robot.data.body_vel_w[:, ee_frame_idx, :] # Extract end-effector velocity in the world frame - root_vel_w = robot.data.root_vel_w # Extract root velocity in the world frame + ee_vel_w = robot.data.body_com_vel_w[:, ee_frame_idx, :] # Extract end-effector velocity in the world frame + root_vel_w = robot.data.root_com_vel_w # 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_rotate_inverse(robot.data.root_quat_w, relative_vel_w[:, 0:3]) # From world to root frame - ee_ang_vel_b = quat_rotate_inverse(robot.data.root_quat_w, relative_vel_w[:, 3:6]) + ee_lin_vel_b = quat_rotate_inverse( + robot.data.root_link_quat_w, relative_vel_w[:, 0:3] + ) # From world to root frame + ee_ang_vel_b = quat_rotate_inverse(robot.data.root_link_quat_w, relative_vel_w[:, 3:6]) ee_vel_b = torch.cat([ee_lin_vel_b, ee_ang_vel_b], dim=-1) # Calculate the contact force diff --git a/source/extensions/omni.isaac.lab/test/envs/check_manager_based_env_floating_cube.py b/source/extensions/omni.isaac.lab/test/envs/check_manager_based_env_floating_cube.py index 4282eb5e8c..812bb98f37 100644 --- a/source/extensions/omni.isaac.lab/test/envs/check_manager_based_env_floating_cube.py +++ b/source/extensions/omni.isaac.lab/test/envs/check_manager_based_env_floating_cube.py @@ -128,11 +128,11 @@ 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 - (self._asset.data.root_pos_w - self._env.scene.env_origins) - vel_error = -self._asset.data.root_lin_vel_w + pos_error = self._processed_actions - (self._asset.data.root_link_pos_w - self._env.scene.env_origins) + vel_error = -self._asset.data.root_com_lin_vel_w # 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) + self._asset.write_root_com_velocity_to_sim(self._vel_command) @configclass @@ -151,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 asset.data.root_pos_w - env.scene.env_origins + return asset.data.root_link_pos_w - env.scene.env_origins ## diff --git a/source/extensions/omni.isaac.lab/test/markers/check_markers_visibility.py b/source/extensions/omni.isaac.lab/test/markers/check_markers_visibility.py index 5aafc68c80..ef15ca5844 100644 --- a/source/extensions/omni.isaac.lab/test/markers/check_markers_visibility.py +++ b/source/extensions/omni.isaac.lab/test/markers/check_markers_visibility.py @@ -98,7 +98,8 @@ def run_simulator( # root state root_state = scene["robot"].data.default_root_state.clone() root_state[:, :3] += scene.env_origins - scene["robot"].write_root_state_to_sim(root_state) + scene["robot"].write_root_link_pose_to_sim(root_state[:, :7]) + scene["robot"].write_root_com_velocity_to_sim(root_state[:, 7:]) # set joint positions with some noise joint_pos, joint_vel = ( scene["robot"].data.default_joint_pos.clone(), diff --git a/source/extensions/omni.isaac.lab/test/scene/check_interactive_scene.py b/source/extensions/omni.isaac.lab/test/scene/check_interactive_scene.py index 25b63fc087..8120102d66 100644 --- a/source/extensions/omni.isaac.lab/test/scene/check_interactive_scene.py +++ b/source/extensions/omni.isaac.lab/test/scene/check_interactive_scene.py @@ -130,11 +130,13 @@ def main(): joint_vel = scene.articulations["robot_1"].data.default_joint_vel # -- set root state # -- robot 1 - scene.articulations["robot_1"].write_root_state_to_sim(root_state) + scene.articulations["robot_1"].write_root_link_pose_to_sim(root_state[:, :7]) + scene.articulations["robot_1"].write_root_com_velocity_to_sim(root_state[:, 7:]) scene.articulations["robot_1"].write_joint_state_to_sim(joint_pos, joint_vel) # -- robot 2 root_state[:, 1] += 1.0 - scene.articulations["robot_2"].write_root_state_to_sim(root_state) + scene.articulations["robot_2"].write_root_link_pose_to_sim(root_state[:, :7]) + scene.articulations["robot_2"].write_root_com_velocity_to_sim(root_state[:, 7:]) scene.articulations["robot_2"].write_joint_state_to_sim(joint_pos, joint_vel) # reset buffers scene.reset() diff --git a/source/extensions/omni.isaac.lab/test/sensors/check_imu_sensor.py b/source/extensions/omni.isaac.lab/test/sensors/check_imu_sensor.py index 6516c0eefb..69a0422455 100644 --- a/source/extensions/omni.isaac.lab/test/sensors/check_imu_sensor.py +++ b/source/extensions/omni.isaac.lab/test/sensors/check_imu_sensor.py @@ -151,8 +151,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_link_pos_w.clone() + ball_initial_orientations = balls.data.root_link_quat_w.clone() # Create a counter for resetting the scene step_count = 0 @@ -168,7 +168,7 @@ def main(): # Reset the scene if step_count % 500 == 0: # reset ball positions - balls.write_root_pose_to_sim(torch.cat([ball_initial_positions, ball_initial_orientations], dim=-1)) + balls.write_root_link_pose_to_sim(torch.cat([ball_initial_positions, ball_initial_orientations], dim=-1)) balls.reset() # reset the sensor imu.reset() diff --git a/source/extensions/omni.isaac.lab/test/sensors/test_contact_sensor.py b/source/extensions/omni.isaac.lab/test/sensors/test_contact_sensor.py index 6bba29b858..949456c3f4 100644 --- a/source/extensions/omni.isaac.lab/test/sensors/test_contact_sensor.py +++ b/source/extensions/omni.isaac.lab/test/sensors/test_contact_sensor.py @@ -405,7 +405,7 @@ def _test_sensor_contact(self, shape: RigidObject, sensor: ContactSensor, mode: duration = self.durations[idx] while current_test_time < duration: # set object states to contact the ground plane - shape.write_root_pose_to_sim(root_pose=test_pose) + shape.write_root_link_pose_to_sim(root_pose=test_pose) # perform simulation step self._perform_sim_step() # increment contact time @@ -432,7 +432,7 @@ def _test_sensor_contact(self, shape: RigidObject, sensor: ContactSensor, mode: dt=duration + self.sim_dt, ) # switch the contact mode for 1 dt step before the next contact test begins. - shape.write_root_pose_to_sim(root_pose=reset_pose) + shape.write_root_link_pose_to_sim(root_pose=reset_pose) # perform simulation step self._perform_sim_step() # set the last air time to 2 sim_dt steps, because last_air_time and last_contact_time diff --git a/source/extensions/omni.isaac.lab/test/sensors/test_frame_transformer.py b/source/extensions/omni.isaac.lab/test/sensors/test_frame_transformer.py index 34321e76da..bc30bed2f3 100644 --- a/source/extensions/omni.isaac.lab/test/sensors/test_frame_transformer.py +++ b/source/extensions/omni.isaac.lab/test/sensors/test_frame_transformer.py @@ -175,7 +175,8 @@ def test_frame_transformer_feet_wrt_base(self): joint_vel = scene.articulations["robot"].data.default_joint_vel # -- set root state # -- robot - scene.articulations["robot"].write_root_state_to_sim(root_state) + scene.articulations["robot"].write_root_link_pose_to_sim(root_state[:, :7]) + scene.articulations["robot"].write_root_com_velocity_to_sim(root_state[:, 7:]) scene.articulations["robot"].write_joint_state_to_sim(joint_pos, joint_vel) # reset buffers scene.reset() @@ -192,9 +193,9 @@ def test_frame_transformer_feet_wrt_base(self): # check absolute frame transforms in world frame # -- ground-truth - root_pose_w = scene.articulations["robot"].data.root_state_w[:, :7] - 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] + root_pose_w = scene.articulations["robot"].data.root_link_state_w[:, :7] + feet_pos_w_gt = scene.articulations["robot"].data.body_link_pos_w[:, feet_indices] + feet_quat_w_gt = scene.articulations["robot"].data.body_link_quat_w[:, feet_indices] # -- 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 @@ -275,7 +276,8 @@ def test_frame_transformer_feet_wrt_thigh(self): joint_vel = scene.articulations["robot"].data.default_joint_vel # -- set root state # -- robot - scene.articulations["robot"].write_root_state_to_sim(root_state) + scene.articulations["robot"].write_root_link_pose_to_sim(root_state[:, :7]) + scene.articulations["robot"].write_root_com_velocity_to_sim(root_state[:, 7:]) scene.articulations["robot"].write_joint_state_to_sim(joint_pos, joint_vel) # reset buffers scene.reset() @@ -292,9 +294,9 @@ def test_frame_transformer_feet_wrt_thigh(self): # check absolute frame transforms in world frame # -- ground-truth - source_pose_w_gt = scene.articulations["robot"].data.body_state_w[:, source_frame_index, :7] - 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] + source_pose_w_gt = scene.articulations["robot"].data.body_link_state_w[:, source_frame_index, :7] + feet_pos_w_gt = scene.articulations["robot"].data.body_link_pos_w[:, feet_indices] + feet_quat_w_gt = scene.articulations["robot"].data.body_link_quat_w[:, feet_indices] # -- 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 @@ -356,7 +358,8 @@ def test_frame_transformer_robot_body_to_external_cube(self): joint_vel = scene.articulations["robot"].data.default_joint_vel # -- set root state # -- robot - scene.articulations["robot"].write_root_state_to_sim(root_state) + scene.articulations["robot"].write_root_link_pose_to_sim(root_state[:, :7]) + scene.articulations["robot"].write_root_com_velocity_to_sim(root_state[:, 7:]) scene.articulations["robot"].write_joint_state_to_sim(joint_pos, joint_vel) # reset buffers scene.reset() @@ -373,9 +376,9 @@ def test_frame_transformer_robot_body_to_external_cube(self): # 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_link_state_w[:, :7] + cube_pos_w_gt = scene.rigid_objects["cube"].data.root_link_state_w[:, :3] + cube_quat_w_gt = scene.rigid_objects["cube"].data.root_link_state_w[:, 3:7] # -- 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 @@ -447,7 +450,8 @@ def test_frame_transformer_offset_frames(self): root_state[:, :3] += scene.env_origins # -- set root state # -- cube - scene["cube"].write_root_state_to_sim(root_state) + scene["cube"].write_root_link_pose_to_sim(root_state[:, :7]) + scene["cube"].write_root_com_velocity_to_sim(root_state[:, 7:]) # reset buffers scene.reset() @@ -460,8 +464,8 @@ def test_frame_transformer_offset_frames(self): # 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_link_state_w[:, :3] + cube_quat_w_gt = scene["cube"].data.root_link_state_w[:, 3:7] # -- 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 @@ -534,7 +538,8 @@ def test_frame_transformer_all_bodies(self): joint_vel = scene.articulations["robot"].data.default_joint_vel # -- set root state # -- robot - scene.articulations["robot"].write_root_state_to_sim(root_state) + scene.articulations["robot"].write_root_link_pose_to_sim(root_state[:, :7]) + scene.articulations["robot"].write_root_com_velocity_to_sim(root_state[:, 7:]) scene.articulations["robot"].write_joint_state_to_sim(joint_pos, joint_vel) # reset buffers scene.reset() @@ -551,9 +556,9 @@ def test_frame_transformer_all_bodies(self): # check absolute frame transforms in world frame # -- ground-truth - root_pose_w = scene.articulations["robot"].data.root_state_w[:, :7] - bodies_pos_w_gt = scene.articulations["robot"].data.body_pos_w - bodies_quat_w_gt = scene.articulations["robot"].data.body_quat_w + root_pose_w = scene.articulations["robot"].data.root_link_state_w[:, :7] + bodies_pos_w_gt = scene.articulations["robot"].data.body_link_pos_w + bodies_quat_w_gt = scene.articulations["robot"].data.body_link_quat_w # -- frame transformer source_pos_w_tf = scene.sensors["frame_transformer"].data.source_pos_w diff --git a/source/extensions/omni.isaac.lab/test/sensors/test_imu.py b/source/extensions/omni.isaac.lab/test/sensors/test_imu.py index e04211d1a9..d7b755dedd 100644 --- a/source/extensions/omni.isaac.lab/test/sensors/test_imu.py +++ b/source/extensions/omni.isaac.lab/test/sensors/test_imu.py @@ -189,12 +189,12 @@ def test_constant_velocity(self): for idx in range(200): # set velocity - self.scene.rigid_objects["balls"].write_root_velocity_to_sim( + self.scene.rigid_objects["balls"].write_root_com_velocity_to_sim( torch.tensor([[1.0, 0.0, 0.0, 0.0, 0.0, 0.0]], dtype=torch.float32, device=self.scene.device).repeat( self.scene.num_envs, 1 ) ) - self.scene.rigid_objects["cube"].write_root_velocity_to_sim( + self.scene.rigid_objects["cube"].write_root_com_velocity_to_sim( torch.tensor([[1.0, 0.0, 0.0, 0.0, 0.0, 0.0]], dtype=torch.float32, device=self.scene.device).repeat( self.scene.num_envs, 1 ) @@ -236,7 +236,7 @@ def test_constant_velocity(self): ) # check the imu velocities - # NOTE: the expected lin_vel_b is the same as the set velocity, as write_root_velocity_to_sim is + # NOTE: the expected lin_vel_b is the same as the set velocity, as write_root_com_velocity_to_sim is # 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( @@ -266,7 +266,7 @@ def test_constant_acceleration(self): """Test the Imu sensor with a constant acceleration.""" for idx in range(100): # set acceleration - self.scene.rigid_objects["balls"].write_root_velocity_to_sim( + self.scene.rigid_objects["balls"].write_root_com_velocity_to_sim( torch.tensor([[0.1, 0.0, 0.0, 0.0, 0.0, 0.0]], dtype=torch.float32, device=self.scene.device).repeat( self.scene.num_envs, 1 ) @@ -287,7 +287,7 @@ def test_constant_acceleration(self): torch.testing.assert_close( self.scene.sensors["imu_ball"].data.lin_acc_b, math_utils.quat_rotate_inverse( - self.scene.rigid_objects["balls"].data.root_quat_w, + self.scene.rigid_objects["balls"].data.root_link_quat_w, torch.tensor([[0.1, 0.0, 0.0]], dtype=torch.float32, device=self.scene.device).repeat( self.scene.num_envs, 1 ) @@ -300,7 +300,7 @@ def test_constant_acceleration(self): # check the angular velocity torch.testing.assert_close( self.scene.sensors["imu_ball"].data.ang_vel_b, - self.scene.rigid_objects["balls"].data.root_ang_vel_b, + self.scene.rigid_objects["balls"].data.root_com_ang_vel_b, rtol=1e-4, atol=1e-4, ) @@ -438,7 +438,7 @@ def test_offset_calculation(self): # should achieve same results between the two imu sensors on the robot for idx in range(500): # set acceleration - self.scene.articulations["robot"].write_root_velocity_to_sim( + self.scene.articulations["robot"].write_root_com_velocity_to_sim( torch.tensor([[0.05, 0.0, 0.0, 0.0, 0.0, 0.0]], dtype=torch.float32, device=self.scene.device).repeat( self.scene.num_envs, 1 ) @@ -504,7 +504,7 @@ def test_env_ids_propogation(self): for idx in range(10): # set acceleration - self.scene.articulations["robot"].write_root_velocity_to_sim( + self.scene.articulations["robot"].write_root_com_velocity_to_sim( torch.tensor([[0.5, 0.0, 0.0, 0.0, 0.0, 0.0]], dtype=torch.float32, device=self.scene.device).repeat( self.scene.num_envs, 1 ) diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/anymal_c/anymal_c_env.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/anymal_c/anymal_c_env.py index 1b2ee577c7..0d2c9f16fd 100644 --- a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/anymal_c/anymal_c_env.py +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/anymal_c/anymal_c_env.py @@ -89,8 +89,8 @@ def _get_observations(self) -> dict: [ tensor for tensor in ( - self._robot.data.root_lin_vel_b, - self._robot.data.root_ang_vel_b, + self._robot.data.root_com_lin_vel_b, + self._robot.data.root_com_ang_vel_b, self._robot.data.projected_gravity_b, self._commands, self._robot.data.joint_pos - self._robot.data.default_joint_pos, @@ -107,15 +107,17 @@ def _get_observations(self) -> dict: def _get_rewards(self) -> torch.Tensor: # linear velocity tracking - lin_vel_error = torch.sum(torch.square(self._commands[:, :2] - self._robot.data.root_lin_vel_b[:, :2]), dim=1) + lin_vel_error = torch.sum( + torch.square(self._commands[:, :2] - self._robot.data.root_com_lin_vel_b[:, :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] - self._robot.data.root_ang_vel_b[:, 2]) + yaw_rate_error = torch.square(self._commands[:, 2] - self._robot.data.root_com_ang_vel_b[:, 2]) yaw_rate_error_mapped = torch.exp(-yaw_rate_error / 0.25) # z velocity tracking - z_vel_error = torch.square(self._robot.data.root_lin_vel_b[:, 2]) + z_vel_error = torch.square(self._robot.data.root_com_lin_vel_b[:, 2]) # angular velocity x/y - ang_vel_error = torch.sum(torch.square(self._robot.data.root_ang_vel_b[:, :2]), dim=1) + ang_vel_error = torch.sum(torch.square(self._robot.data.root_com_ang_vel_b[:, :2]), dim=1) # joint torques joint_torques = torch.sum(torch.square(self._robot.data.applied_torque), dim=1) # joint acceleration @@ -178,8 +180,8 @@ def _reset_idx(self, env_ids: torch.Tensor | None): joint_vel = self._robot.data.default_joint_vel[env_ids] default_root_state = self._robot.data.default_root_state[env_ids] default_root_state[:, :3] += self._terrain.env_origins[env_ids] - self._robot.write_root_pose_to_sim(default_root_state[:, :7], env_ids) - self._robot.write_root_velocity_to_sim(default_root_state[:, 7:], env_ids) + self._robot.write_root_link_pose_to_sim(default_root_state[:, :7], env_ids) + self._robot.write_root_com_velocity_to_sim(default_root_state[:, 7:], env_ids) self._robot.write_joint_state_to_sim(joint_pos, joint_vel, None, env_ids) # Logging extras = dict() diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/cart_double_pendulum/cart_double_pendulum_env.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/cart_double_pendulum/cart_double_pendulum_env.py index ad8c616940..25f9e59c79 100644 --- a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/cart_double_pendulum/cart_double_pendulum_env.py +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/cart_double_pendulum/cart_double_pendulum_env.py @@ -182,8 +182,8 @@ def _reset_idx(self, env_ids: Sequence[int] | None): self.joint_pos[env_ids] = joint_pos self.joint_vel[env_ids] = joint_vel - self.robot.write_root_pose_to_sim(default_root_state[:, :7], env_ids) - self.robot.write_root_velocity_to_sim(default_root_state[:, 7:], env_ids) + self.robot.write_root_link_pose_to_sim(default_root_state[:, :7], env_ids) + self.robot.write_root_com_velocity_to_sim(default_root_state[:, 7:], env_ids) self.robot.write_joint_state_to_sim(joint_pos, joint_vel, None, env_ids) diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/cartpole/cartpole_camera_env.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/cartpole/cartpole_camera_env.py index e5881db086..fea5f5511a 100644 --- a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/cartpole/cartpole_camera_env.py +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/cartpole/cartpole_camera_env.py @@ -199,8 +199,8 @@ def _reset_idx(self, env_ids: Sequence[int] | None): self.joint_pos[env_ids] = joint_pos self.joint_vel[env_ids] = joint_vel - self._cartpole.write_root_pose_to_sim(default_root_state[:, :7], env_ids) - self._cartpole.write_root_velocity_to_sim(default_root_state[:, 7:], env_ids) + self._cartpole.write_root_link_pose_to_sim(default_root_state[:, :7], env_ids) + self._cartpole.write_root_com_velocity_to_sim(default_root_state[:, 7:], env_ids) self._cartpole.write_joint_state_to_sim(joint_pos, joint_vel, None, env_ids) diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/cartpole/cartpole_env.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/cartpole/cartpole_env.py index 1525941672..88b1877d08 100644 --- a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/cartpole/cartpole_env.py +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/cartpole/cartpole_env.py @@ -143,8 +143,8 @@ def _reset_idx(self, env_ids: Sequence[int] | None): self.joint_pos[env_ids] = joint_pos self.joint_vel[env_ids] = joint_vel - self.cartpole.write_root_pose_to_sim(default_root_state[:, :7], env_ids) - self.cartpole.write_root_velocity_to_sim(default_root_state[:, 7:], env_ids) + self.cartpole.write_root_link_pose_to_sim(default_root_state[:, :7], env_ids) + self.cartpole.write_root_com_velocity_to_sim(default_root_state[:, 7:], env_ids) self.cartpole.write_joint_state_to_sim(joint_pos, joint_vel, None, env_ids) diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/factory/factory_env.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/factory/factory_env.py index df7850d32c..42af6e531c 100644 --- a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/factory/factory_env.py +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/factory/factory_env.py @@ -189,16 +189,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 = self._fixed_asset.data.root_pos_w - self.scene.env_origins - self.fixed_quat = self._fixed_asset.data.root_quat_w + self.fixed_pos = self._fixed_asset.data.root_link_pos_w - self.scene.env_origins + self.fixed_quat = self._fixed_asset.data.root_link_quat_w - self.held_pos = self._held_asset.data.root_pos_w - self.scene.env_origins - self.held_quat = self._held_asset.data.root_quat_w + self.held_pos = self._held_asset.data.root_link_pos_w - self.scene.env_origins + self.held_quat = self._held_asset.data.root_link_quat_w - self.fingertip_midpoint_pos = self._robot.data.body_pos_w[:, self.fingertip_body_idx] - self.scene.env_origins - self.fingertip_midpoint_quat = self._robot.data.body_quat_w[:, self.fingertip_body_idx] - self.fingertip_midpoint_linvel = self._robot.data.body_lin_vel_w[:, self.fingertip_body_idx] - self.fingertip_midpoint_angvel = self._robot.data.body_ang_vel_w[:, self.fingertip_body_idx] + self.fingertip_midpoint_pos = ( + self._robot.data.body_link_pos_w[:, self.fingertip_body_idx] - self.scene.env_origins + ) + self.fingertip_midpoint_quat = self._robot.data.body_link_quat_w[:, self.fingertip_body_idx] + self.fingertip_midpoint_linvel = self._robot.data.body_com_lin_vel_w[:, self.fingertip_body_idx] + self.fingertip_midpoint_angvel = self._robot.data.body_com_ang_vel_w[:, self.fingertip_body_idx] jacobians = self._robot.root_physx_view.get_jacobians() @@ -552,13 +554,15 @@ def _set_assets_to_default_pose(self, env_ids): held_state = self._held_asset.data.default_root_state.clone()[env_ids] held_state[:, 0:3] += self.scene.env_origins[env_ids] held_state[:, 7:] = 0.0 - self._held_asset.write_root_state_to_sim(held_state, env_ids=env_ids) + self._held_asset.write_root_link_pose_to_sim(held_state[:, 0:7], env_ids=env_ids) + self._held_asset.write_root_com_velocity_to_sim(held_state[:, 7:], env_ids=env_ids) self._held_asset.reset() fixed_state = self._fixed_asset.data.default_root_state.clone()[env_ids] fixed_state[:, 0:3] += self.scene.env_origins[env_ids] fixed_state[:, 7:] = 0.0 - self._fixed_asset.write_root_state_to_sim(fixed_state, env_ids=env_ids) + self._fixed_asset.write_root_link_pose_to_sim(fixed_state[:, 0:7], env_ids=env_ids) + self._fixed_asset.write_root_com_velocity_to_sim(fixed_state[:, 7:], env_ids=env_ids) self._fixed_asset.reset() def set_pos_inverse_kinematics(self, env_ids): @@ -681,7 +685,8 @@ def randomize_initial_state(self, env_ids): # (1.c.) Velocity fixed_state[:, 7:] = 0.0 # vel # (1.d.) Update values. - self._fixed_asset.write_root_state_to_sim(fixed_state, env_ids=env_ids) + self._fixed_asset.write_root_link_pose_to_sim(fixed_state[:, 0:7], env_ids=env_ids) + self._fixed_asset.write_root_com_velocity_to_sim(fixed_state[:, 7:], env_ids=env_ids) self._fixed_asset.reset() # (1.e.) Noisy position observation. @@ -767,13 +772,15 @@ def randomize_initial_state(self, env_ids): small_gear_state = self._small_gear_asset.data.default_root_state.clone()[env_ids] small_gear_state[:, 0:7] = fixed_state[:, 0:7] small_gear_state[:, 7:] = 0.0 # vel - self._small_gear_asset.write_root_state_to_sim(small_gear_state, env_ids=env_ids) + self._small_gear_asset.write_root_link_pose_to_sim(small_gear_state[:, 0:7], env_ids=env_ids) + self._small_gear_asset.write_root_com_velocity_to_sim(small_gear_state[:, 7:], env_ids=env_ids) self._small_gear_asset.reset() large_gear_state = self._large_gear_asset.data.default_root_state.clone()[env_ids] large_gear_state[:, 0:7] = fixed_state[:, 0:7] large_gear_state[:, 7:] = 0.0 # vel - self._large_gear_asset.write_root_state_to_sim(large_gear_state, env_ids=env_ids) + self._large_gear_asset.write_root_link_pose_to_sim(large_gear_state[:, 0:7], env_ids=env_ids) + self._large_gear_asset.write_root_com_velocity_to_sim(large_gear_state[:, 7:], env_ids=env_ids) self._large_gear_asset.reset() # (3) Randomize asset-in-gripper location. @@ -815,7 +822,8 @@ def randomize_initial_state(self, env_ids): held_state[:, 0:3] = translated_held_asset_pos + self.scene.env_origins held_state[:, 3:7] = translated_held_asset_quat held_state[:, 7:] = 0.0 - self._held_asset.write_root_state_to_sim(held_state) + self._held_asset.write_root_link_pose_to_sim(held_state[:, 0:7]) + self._held_asset.write_root_com_velocity_to_sim(held_state[:, 7:]) self._held_asset.reset() # Close hand diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/franka_cabinet/franka_cabinet_env.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/franka_cabinet/franka_cabinet_env.py index 3a6a480ed0..d7234ca386 100644 --- a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/franka_cabinet/franka_cabinet_env.py +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/franka_cabinet/franka_cabinet_env.py @@ -299,8 +299,8 @@ def _get_dones(self) -> tuple[torch.Tensor, torch.Tensor]: def _get_rewards(self) -> torch.Tensor: # Refresh the intermediate values after the physics steps self._compute_intermediate_values() - robot_left_finger_pos = self._robot.data.body_pos_w[:, self.left_finger_link_idx] - robot_right_finger_pos = self._robot.data.body_pos_w[:, self.right_finger_link_idx] + robot_left_finger_pos = self._robot.data.body_link_pos_w[:, self.left_finger_link_idx] + robot_right_finger_pos = self._robot.data.body_link_pos_w[:, self.right_finger_link_idx] return self._compute_rewards( self.actions, @@ -372,10 +372,10 @@ def _compute_intermediate_values(self, env_ids: torch.Tensor | None = None): if env_ids is None: env_ids = self._robot._ALL_INDICES - hand_pos = self._robot.data.body_pos_w[env_ids, self.hand_link_idx] - hand_rot = self._robot.data.body_quat_w[env_ids, self.hand_link_idx] - drawer_pos = self._cabinet.data.body_pos_w[env_ids, self.drawer_link_idx] - drawer_rot = self._cabinet.data.body_quat_w[env_ids, self.drawer_link_idx] + hand_pos = self._robot.data.body_link_pos_w[env_ids, self.hand_link_idx] + hand_rot = self._robot.data.body_link_quat_w[env_ids, self.hand_link_idx] + drawer_pos = self._cabinet.data.body_link_pos_w[env_ids, self.drawer_link_idx] + drawer_rot = self._cabinet.data.body_link_quat_w[env_ids, self.drawer_link_idx] ( self.robot_grasp_rot[env_ids], self.robot_grasp_pos[env_ids], diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/inhand_manipulation/inhand_manipulation_env.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/inhand_manipulation/inhand_manipulation_env.py index dd671d2ab9..8017889c6d 100644 --- a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/inhand_manipulation/inhand_manipulation_env.py +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/inhand_manipulation/inhand_manipulation_env.py @@ -221,7 +221,8 @@ def _reset_idx(self, env_ids: Sequence[int] | None): ) object_default_state[:, 7:] = torch.zeros_like(self.object.data.default_root_state[env_ids, 7:]) - self.object.write_root_state_to_sim(object_default_state, env_ids) + self.object.write_root_link_pose_to_sim(object_default_state[:, :7], env_ids) + self.object.write_root_com_velocity_to_sim(object_default_state[:, 7:], env_ids) # reset hand delta_max = self.hand_dof_upper_limits[env_ids] - self.hand.data.default_joint_pos[env_ids] @@ -260,22 +261,22 @@ def _reset_target_pose(self, env_ids): def _compute_intermediate_values(self): # data for hand - self.fingertip_pos = self.hand.data.body_pos_w[:, self.finger_bodies] - self.fingertip_rot = self.hand.data.body_quat_w[:, self.finger_bodies] + self.fingertip_pos = self.hand.data.body_link_pos_w[:, self.finger_bodies] + self.fingertip_rot = self.hand.data.body_link_quat_w[:, 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 = self.hand.data.body_vel_w[:, self.finger_bodies] + self.fingertip_velocities = self.hand.data.body_com_vel_w[:, self.finger_bodies] self.hand_dof_pos = self.hand.data.joint_pos self.hand_dof_vel = self.hand.data.joint_vel # data for object - self.object_pos = self.object.data.root_pos_w - self.scene.env_origins - self.object_rot = self.object.data.root_quat_w - self.object_velocities = self.object.data.root_vel_w - self.object_linvel = self.object.data.root_lin_vel_w - self.object_angvel = self.object.data.root_ang_vel_w + self.object_pos = self.object.data.root_link_pos_w - self.scene.env_origins + self.object_rot = self.object.data.root_link_quat_w + self.object_velocities = self.object.data.root_com_vel_w + self.object_linvel = self.object.data.root_com_lin_vel_w + self.object_angvel = self.object.data.root_com_ang_vel_w def compute_reduced_observations(self): # Per https://arxiv.org/pdf/1808.00177.pdf Table 2 diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/locomotion/locomotion_env.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/locomotion/locomotion_env.py index ffa572ac41..dc3a5439da 100644 --- a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/locomotion/locomotion_env.py +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/locomotion/locomotion_env.py @@ -68,8 +68,8 @@ def _apply_action(self): self.robot.set_joint_effort_target(forces, joint_ids=self._joint_dof_idx) def _compute_intermediate_values(self): - self.torso_position, self.torso_rotation = self.robot.data.root_pos_w, self.robot.data.root_quat_w - self.velocity, self.ang_velocity = self.robot.data.root_lin_vel_w, self.robot.data.root_ang_vel_w + self.torso_position, self.torso_rotation = self.robot.data.root_link_pos_w, self.robot.data.root_link_quat_w + self.velocity, self.ang_velocity = self.robot.data.root_com_lin_vel_w, self.robot.data.root_com_ang_vel_w self.dof_pos, self.dof_vel = self.robot.data.joint_pos, self.robot.data.joint_vel ( @@ -161,8 +161,8 @@ def _reset_idx(self, env_ids: torch.Tensor | None): default_root_state = self.robot.data.default_root_state[env_ids] default_root_state[:, :3] += self.scene.env_origins[env_ids] - self.robot.write_root_pose_to_sim(default_root_state[:, :7], env_ids) - self.robot.write_root_velocity_to_sim(default_root_state[:, 7:], env_ids) + self.robot.write_root_link_pose_to_sim(default_root_state[:, :7], env_ids) + self.robot.write_root_com_velocity_to_sim(default_root_state[:, 7:], env_ids) self.robot.write_joint_state_to_sim(joint_pos, joint_vel, None, env_ids) to_target = self.targets[env_ids] - default_root_state[:, :3] diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/quadcopter/quadcopter_env.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/quadcopter/quadcopter_env.py index 97156618f1..7389975ebc 100644 --- a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/quadcopter/quadcopter_env.py +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/quadcopter/quadcopter_env.py @@ -154,12 +154,12 @@ def _apply_action(self): 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_link_state_w[:, :3], self._robot.data.root_link_state_w[:, 3:7], self._desired_pos_w ) obs = torch.cat( [ - self._robot.data.root_lin_vel_b, - self._robot.data.root_ang_vel_b, + self._robot.data.root_com_lin_vel_b, + self._robot.data.root_com_ang_vel_b, self._robot.data.projected_gravity_b, desired_pos_b, ], @@ -169,9 +169,9 @@ def _get_observations(self) -> dict: return observations def _get_rewards(self) -> torch.Tensor: - lin_vel = torch.sum(torch.square(self._robot.data.root_lin_vel_b), dim=1) - ang_vel = torch.sum(torch.square(self._robot.data.root_ang_vel_b), dim=1) - distance_to_goal = torch.linalg.norm(self._desired_pos_w - self._robot.data.root_pos_w, dim=1) + lin_vel = torch.sum(torch.square(self._robot.data.root_com_lin_vel_b), dim=1) + ang_vel = torch.sum(torch.square(self._robot.data.root_com_ang_vel_b), dim=1) + distance_to_goal = torch.linalg.norm(self._desired_pos_w - self._robot.data.root_link_pos_w, 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, @@ -186,7 +186,9 @@ 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(self._robot.data.root_pos_w[:, 2] < 0.1, self._robot.data.root_pos_w[:, 2] > 2.0) + died = torch.logical_or( + self._robot.data.root_link_pos_w[:, 2] < 0.1, self._robot.data.root_link_pos_w[:, 2] > 2.0 + ) return died, time_out def _reset_idx(self, env_ids: torch.Tensor | None): @@ -195,7 +197,7 @@ def _reset_idx(self, env_ids: torch.Tensor | None): # Logging final_distance_to_goal = torch.linalg.norm( - self._desired_pos_w[env_ids] - self._robot.data.root_pos_w[env_ids], dim=1 + self._desired_pos_w[env_ids] - self._robot.data.root_link_pos_w[env_ids], dim=1 ).mean() extras = dict() for key in self._episode_sums.keys(): @@ -226,8 +228,8 @@ def _reset_idx(self, env_ids: torch.Tensor | None): joint_vel = self._robot.data.default_joint_vel[env_ids] default_root_state = self._robot.data.default_root_state[env_ids] default_root_state[:, :3] += self._terrain.env_origins[env_ids] - self._robot.write_root_pose_to_sim(default_root_state[:, :7], env_ids) - self._robot.write_root_velocity_to_sim(default_root_state[:, 7:], env_ids) + self._robot.write_root_link_pose_to_sim(default_root_state[:, :7], env_ids) + self._robot.write_root_com_velocity_to_sim(default_root_state[:, 7:], env_ids) self._robot.write_joint_state_to_sim(joint_pos, joint_vel, None, env_ids) def _set_debug_vis_impl(self, debug_vis: bool): diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/shadow_hand_over/shadow_hand_over_env.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/shadow_hand_over/shadow_hand_over_env.py index 01ac2c22d7..622a95e842 100644 --- a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/shadow_hand_over/shadow_hand_over_env.py +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/shadow_hand_over/shadow_hand_over_env.py @@ -322,7 +322,8 @@ def _reset_idx(self, env_ids: Sequence[int] | torch.Tensor | None): ) object_default_state[:, 7:] = torch.zeros_like(self.object.data.default_root_state[env_ids, 7:]) - self.object.write_root_state_to_sim(object_default_state, env_ids) + self.object.write_root_link_pose_to_sim(object_default_state[:, :7], env_ids) + self.object.write_root_com_velocity_to_sim(object_default_state[:, 7:], env_ids) # reset right hand delta_max = self.hand_dof_upper_limits[env_ids] - self.right_hand.data.default_joint_pos[env_ids] @@ -376,33 +377,33 @@ def _reset_target_pose(self, env_ids): def _compute_intermediate_values(self): # data for right hand - self.right_fingertip_pos = self.right_hand.data.body_pos_w[:, self.finger_bodies] - self.right_fingertip_rot = self.right_hand.data.body_quat_w[:, self.finger_bodies] + self.right_fingertip_pos = self.right_hand.data.body_link_pos_w[:, self.finger_bodies] + self.right_fingertip_rot = self.right_hand.data.body_link_quat_w[:, 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 = self.right_hand.data.body_vel_w[:, self.finger_bodies] + self.right_fingertip_velocities = self.right_hand.data.body_com_vel_w[:, self.finger_bodies] self.right_hand_dof_pos = self.right_hand.data.joint_pos self.right_hand_dof_vel = self.right_hand.data.joint_vel # data for left hand - self.left_fingertip_pos = self.left_hand.data.body_pos_w[:, self.finger_bodies] - self.left_fingertip_rot = self.left_hand.data.body_quat_w[:, self.finger_bodies] + self.left_fingertip_pos = self.left_hand.data.body_link_pos_w[:, self.finger_bodies] + self.left_fingertip_rot = self.left_hand.data.body_link_quat_w[:, 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 = self.left_hand.data.body_vel_w[:, self.finger_bodies] + self.left_fingertip_velocities = self.left_hand.data.body_com_vel_w[:, self.finger_bodies] self.left_hand_dof_pos = self.left_hand.data.joint_pos self.left_hand_dof_vel = self.left_hand.data.joint_vel # data for object - self.object_pos = self.object.data.root_pos_w - self.scene.env_origins - self.object_rot = self.object.data.root_quat_w - self.object_velocities = self.object.data.root_vel_w - self.object_linvel = self.object.data.root_lin_vel_w - self.object_angvel = self.object.data.root_ang_vel_w + self.object_pos = self.object.data.root_link_pos_w - self.scene.env_origins + self.object_rot = self.object.data.root_link_quat_w + self.object_velocities = self.object.data.root_com_vel_w + self.object_linvel = self.object.data.root_com_lin_vel_w + self.object_angvel = self.object.data.root_com_ang_vel_w @torch.jit.script diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/classic/humanoid/mdp/observations.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/classic/humanoid/mdp/observations.py index 1a65c381a7..e74b872582 100644 --- a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/classic/humanoid/mdp/observations.py +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/classic/humanoid/mdp/observations.py @@ -21,7 +21,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(asset.data.root_quat_w) + roll, _, yaw = math_utils.euler_xyz_from_quat(asset.data.root_link_quat_w) # normalize angle to [-pi, pi] roll = torch.atan2(torch.sin(roll), torch.cos(roll)) yaw = torch.atan2(torch.sin(yaw), torch.cos(yaw)) @@ -46,11 +46,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) - asset.data.root_pos_w[:, :3] + to_target_pos = torch.tensor(target_pos, device=env.device) - asset.data.root_link_pos_w[:, :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_rotate(asset.data.root_quat_w, asset.data.FORWARD_VEC_B) + heading_vec = math_utils.quat_rotate(asset.data.root_link_quat_w, asset.data.FORWARD_VEC_B) # 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)) @@ -64,10 +64,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) - asset.data.root_pos_w[:, :3] + to_target_pos = torch.tensor(target_pos, device=env.device) - asset.data.root_link_pos_w[:, :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(asset.data.root_quat_w) + _, _, yaw = math_utils.euler_xyz_from_quat(asset.data.root_link_quat_w) # 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/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/classic/humanoid/mdp/rewards.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/classic/humanoid/mdp/rewards.py index b21be8d31f..913be8c1d2 100644 --- a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/classic/humanoid/mdp/rewards.py +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/classic/humanoid/mdp/rewards.py @@ -53,7 +53,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 - asset.data.root_pos_w[env_ids, :3] + to_target_pos = target_pos - asset.data.root_link_pos_w[env_ids, :3] # reward terms self.potentials[env_ids] = -torch.norm(to_target_pos, p=2, dim=-1) / self._env.step_dt self.prev_potentials[env_ids] = self.potentials[env_ids] @@ -68,7 +68,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 - asset.data.root_pos_w[:, :3] + to_target_pos = target_pos - asset.data.root_link_pos_w[:, :3] to_target_pos[:, 2] = 0.0 # update history buffer and compute new potential self.prev_potentials[:] = self.potentials[:] diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/locomotion/velocity/config/spot/mdp/rewards.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/locomotion/velocity/config/spot/mdp/rewards.py index d4deab6a6a..f587827f73 100644 --- a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/locomotion/velocity/config/spot/mdp/rewards.py +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/locomotion/velocity/config/spot/mdp/rewards.py @@ -49,7 +49,7 @@ def air_time_reward( 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.norm(env.command_manager.get_command("base_velocity"), dim=1).unsqueeze(dim=1).expand(-1, 4) - body_vel = torch.linalg.norm(asset.data.root_lin_vel_b[:, :2], dim=1).unsqueeze(dim=1).expand(-1, 4) + body_vel = torch.linalg.norm(asset.data.root_com_lin_vel_b[:, :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), @@ -64,7 +64,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 - asset.data.root_ang_vel_b[:, 2]).unsqueeze(1), dim=1) + ang_vel_error = torch.linalg.norm((target - asset.data.root_com_ang_vel_b[:, 2]).unsqueeze(1), dim=1) return torch.exp(-ang_vel_error / std) @@ -76,7 +76,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 - asset.data.root_lin_vel_b[:, :2]), dim=1) + lin_vel_error = torch.linalg.norm((target - asset.data.root_com_lin_vel_b[:, :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) @@ -148,7 +148,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.norm(env.command_manager.get_command("base_velocity"), dim=1) - body_vel = torch.linalg.norm(self.asset.data.root_lin_vel_b[:, :2], dim=1) + body_vel = torch.linalg.norm(self.asset.data.root_com_lin_vel_b[:, :2], dim=1) return torch.where( torch.logical_or(cmd > 0.0, body_vel > self.velocity_threshold), sync_reward * async_reward, 0.0 ) @@ -182,8 +182,10 @@ 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(asset.data.body_pos_w[:, asset_cfg.body_ids, 2] - target_height) - foot_velocity_tanh = torch.tanh(tanh_mult * torch.norm(asset.data.body_lin_vel_w[:, asset_cfg.body_ids, :2], dim=2)) + foot_z_target_error = torch.square(asset.data.body_link_pos_w[:, asset_cfg.body_ids, 2] - target_height) + foot_velocity_tanh = torch.tanh( + tanh_mult * torch.norm(asset.data.body_com_lin_vel_w[:, asset_cfg.body_ids, :2], dim=2) + ) reward = foot_z_target_error * foot_velocity_tanh return torch.exp(-torch.sum(reward, dim=1) / std) @@ -217,8 +219,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(asset.data.root_lin_vel_b[:, 2]) + 0.2 * torch.sum( - torch.abs(asset.data.root_ang_vel_b[:, :2]), dim=1 + return 0.8 * torch.square(asset.data.root_com_lin_vel_b[:, 2]) + 0.2 * torch.sum( + torch.abs(asset.data.root_com_ang_vel_b[:, :2]), dim=1 ) @@ -243,7 +245,7 @@ def foot_slip_penalty( # check if contact force is above threshold net_contact_forces = contact_sensor.data.net_forces_w_history is_contact = torch.max(torch.norm(net_contact_forces[:, :, sensor_cfg.body_ids], dim=-1), dim=1)[0] > threshold - foot_planar_velocity = torch.linalg.norm(asset.data.body_lin_vel_w[:, asset_cfg.body_ids, :2], dim=2) + foot_planar_velocity = torch.linalg.norm(asset.data.body_com_lin_vel_w[:, asset_cfg.body_ids, :2], dim=2) reward = is_contact * foot_planar_velocity return torch.sum(reward, dim=1) @@ -263,7 +265,7 @@ 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(asset.data.root_lin_vel_b[:, :2], dim=1) + body_vel = torch.linalg.norm(asset.data.root_com_lin_vel_b[:, :2], dim=1) reward = torch.linalg.norm((asset.data.joint_pos - asset.data.default_joint_pos), dim=1) return torch.where(torch.logical_or(cmd > 0.0, body_vel > velocity_threshold), reward, stand_still_scale * reward) diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/locomotion/velocity/mdp/curriculums.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/locomotion/velocity/mdp/curriculums.py index 38e7b93d10..8092acca2d 100644 --- a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/locomotion/velocity/mdp/curriculums.py +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/locomotion/velocity/mdp/curriculums.py @@ -43,7 +43,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.norm(asset.data.root_pos_w[env_ids, :2] - env.scene.env_origins[env_ids, :2], dim=1) + distance = torch.norm(asset.data.root_link_pos_w[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/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/locomotion/velocity/mdp/rewards.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/locomotion/velocity/mdp/rewards.py index 08c1df0097..07fa615f99 100644 --- a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/locomotion/velocity/mdp/rewards.py +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/locomotion/velocity/mdp/rewards.py @@ -77,7 +77,8 @@ def feet_slide(env, sensor_cfg: SceneEntityCfg, asset_cfg: SceneEntityCfg = Scen contact_sensor: ContactSensor = env.scene.sensors[sensor_cfg.name] contacts = contact_sensor.data.net_forces_w_history[:, :, sensor_cfg.body_ids, :].norm(dim=-1).max(dim=1)[0] > 1.0 asset = env.scene[asset_cfg.name] - body_vel = asset.data.body_lin_vel_w[:, sensor_cfg.body_ids, :2] + + body_vel = asset.data.body_com_lin_vel_w[:, asset_cfg.body_ids, :2] reward = torch.sum(body_vel.norm(dim=-1) * contacts, dim=1) return reward @@ -88,7 +89,7 @@ def track_lin_vel_xy_yaw_frame_exp( """Reward tracking of linear velocity commands (xy axes) in the gravity aligned robot frame using exponential kernel.""" # extract the used quantities (to enable type-hinting) asset = env.scene[asset_cfg.name] - vel_yaw = quat_rotate_inverse(yaw_quat(asset.data.root_quat_w), asset.data.root_lin_vel_w[:, :3]) + vel_yaw = quat_rotate_inverse(yaw_quat(asset.data.root_link_quat_w), asset.data.root_com_lin_vel_w[:, :3]) lin_vel_error = torch.sum( torch.square(env.command_manager.get_command(command_name)[:, :2] - vel_yaw[:, :2]), dim=1 ) @@ -101,5 +102,7 @@ def track_ang_vel_z_world_exp( """Reward tracking of angular velocity commands (yaw) in world frame using exponential kernel.""" # 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] - asset.data.root_ang_vel_w[:, 2]) + ang_vel_error = torch.square( + env.command_manager.get_command(command_name)[:, 2] - asset.data.root_com_ang_vel_w[:, 2] + ) return torch.exp(-ang_vel_error / std**2) diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/locomotion/velocity/mdp/terminations.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/locomotion/velocity/mdp/terminations.py index 464716deec..cd504f0b6f 100644 --- a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/locomotion/velocity/mdp/terminations.py +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/locomotion/velocity/mdp/terminations.py @@ -45,8 +45,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(asset.data.root_pos_w[:, 0]) > 0.5 * map_width - distance_buffer - y_out_of_bounds = torch.abs(asset.data.root_pos_w[:, 1]) > 0.5 * map_height - distance_buffer + x_out_of_bounds = torch.abs(asset.data.root_link_pos_w[:, 0]) > 0.5 * map_width - distance_buffer + y_out_of_bounds = torch.abs(asset.data.root_link_pos_w[:, 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/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/cabinet/mdp/observations.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/cabinet/mdp/observations.py index 6c13b989cc..18d7df8347 100644 --- a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/cabinet/mdp/observations.py +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/cabinet/mdp/observations.py @@ -21,7 +21,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 object_data.root_pos_w - ee_tf_data.target_pos_w[..., 0, :] + return object_data.root_link_pos_w - ee_tf_data.target_pos_w[..., 0, :] def rel_ee_drawer_distance(env: ManagerBasedRLEnv) -> torch.Tensor: diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/inhand/mdp/commands/orientation_command.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/inhand/mdp/commands/orientation_command.py index 153127c45c..167a0b9a43 100644 --- a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/inhand/mdp/commands/orientation_command.py +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/inhand/mdp/commands/orientation_command.py @@ -95,10 +95,10 @@ def _update_metrics(self): # logs data # -- compute the orientation error self.metrics["orientation_error"] = math_utils.quat_error_magnitude( - self.object.data.root_quat_w, self.quat_command_w + self.object.data.root_link_quat_w, self.quat_command_w ) # -- compute the position error - self.metrics["position_error"] = torch.norm(self.object.data.root_pos_w - self.pos_command_w, dim=1) + self.metrics["position_error"] = torch.norm(self.object.data.root_link_pos_w - self.pos_command_w, dim=1) # -- compute the number of consecutive successes successes = self.metrics["orientation_error"] < self.cfg.orientation_success_threshold self.metrics["consecutive_success"] += successes.float() diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/inhand/mdp/observations.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/inhand/mdp/observations.py index a7d4335842..f23ebc716a 100644 --- a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/inhand/mdp/observations.py +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/inhand/mdp/observations.py @@ -30,7 +30,7 @@ def goal_quat_diff( # obtain the orientations goal_quat_w = command_term.command[:, 3:7] - asset_quat_w = asset.data.root_quat_w + asset_quat_w = asset.data.root_link_quat_w # compute quaternion difference quat = math_utils.quat_mul(asset_quat_w, math_utils.quat_conjugate(goal_quat_w)) diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/inhand/mdp/rewards.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/inhand/mdp/rewards.py index 302b85871c..c02f2e301a 100644 --- a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/inhand/mdp/rewards.py +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/inhand/mdp/rewards.py @@ -39,7 +39,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(asset.data.root_quat_w, goal_quat_w) + dtheta = math_utils.quat_error_magnitude(asset.data.root_link_quat_w, goal_quat_w) return dtheta <= threshold @@ -63,7 +63,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 = asset.data.root_pos_w - env.scene.env_origins + object_pos_e = asset.data.root_link_pos_w - env.scene.env_origins return torch.norm(goal_pos_e - object_pos_e, p=2, dim=-1) @@ -91,6 +91,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(asset.data.root_quat_w, goal_quat_w) + dtheta = math_utils.quat_error_magnitude(asset.data.root_link_quat_w, goal_quat_w) return 1.0 / (dtheta + rot_eps) diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/inhand/mdp/terminations.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/inhand/mdp/terminations.py index a56df2a8b9..4aaea26ad8 100644 --- a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/inhand/mdp/terminations.py +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/inhand/mdp/terminations.py @@ -50,7 +50,7 @@ def object_away_from_goal( asset = env.scene[object_cfg.name] # object pos - asset_pos_e = asset.data.root_pos_w - env.scene.env_origins + asset_pos_e = asset.data.root_link_pos_w - env.scene.env_origins goal_pos_e = command_term.command[:, :3] return torch.norm(asset_pos_e - goal_pos_e, p=2, dim=1) > threshold @@ -78,6 +78,6 @@ def object_away_from_robot( object = env.scene[object_cfg.name] # compute distance - dist = torch.norm(robot.data.root_pos_w - object.data.root_pos_w, dim=1) + dist = torch.norm(robot.data.root_link_pos_w - object.data.root_link_pos_w, dim=1) return dist > threshold diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/lift/mdp/observations.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/lift/mdp/observations.py index 3a976fa3f5..46468e23f9 100644 --- a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/lift/mdp/observations.py +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/lift/mdp/observations.py @@ -24,8 +24,8 @@ 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 = object.data.root_pos_w[:, :3] + object_pos_w = object.data.root_link_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 + robot.data.root_link_state_w[:, :3], robot.data.root_link_state_w[:, 3:7], object_pos_w ) return object_pos_b diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/lift/mdp/rewards.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/lift/mdp/rewards.py index 334df9ea50..9c48814668 100644 --- a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/lift/mdp/rewards.py +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/lift/mdp/rewards.py @@ -22,7 +22,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(object.data.root_pos_w[:, 2] > minimal_height, 1.0, 0.0) + return torch.where(object.data.root_link_pos_w[:, 2] > minimal_height, 1.0, 0.0) def object_ee_distance( @@ -36,7 +36,7 @@ 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 = object.data.root_pos_w + cube_pos_w = object.data.root_link_pos_w # End-effector position: (num_envs, 3) ee_w = ee_frame.data.target_pos_w[..., 0, :] # Distance of the end-effector to the object: (num_envs,) @@ -60,8 +60,10 @@ 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_link_state_w[:, :3], robot.data.root_link_state_w[:, 3:7], 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_link_pos_w[:, :3], 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)) + return (object.data.root_link_pos_w[:, 2] > minimal_height) * (1 - torch.tanh(distance / std)) diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/lift/mdp/terminations.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/lift/mdp/terminations.py index 212e192336..bd86490734 100644 --- a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/lift/mdp/terminations.py +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/lift/mdp/terminations.py @@ -45,9 +45,11 @@ 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_link_state_w[:, :3], robot.data.root_link_state_w[:, 3:7], 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_link_pos_w[:, :3], dim=1) # rewarded if the object is lifted above the threshold return distance < threshold diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/reach/mdp/rewards.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/reach/mdp/rewards.py index d5ae2a57c5..8f93c7f70a 100644 --- a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/reach/mdp/rewards.py +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/reach/mdp/rewards.py @@ -28,8 +28,10 @@ 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_link_state_w[:, :3], asset.data.root_link_state_w[:, 3:7], des_pos_b + ) + curr_pos_w = asset.data.body_link_state_w[:, asset_cfg.body_ids[0], :3] # type: ignore return torch.norm(curr_pos_w - des_pos_w, dim=1) @@ -46,8 +48,10 @@ 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_link_state_w[:, :3], asset.data.root_link_state_w[:, 3:7], des_pos_b + ) + curr_pos_w = asset.data.body_link_state_w[:, asset_cfg.body_ids[0], :3] # type: ignore distance = torch.norm(curr_pos_w - des_pos_w, dim=1) return 1 - torch.tanh(distance / std) @@ -64,6 +68,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_link_state_w[:, 3:7], des_quat_b) + curr_quat_w = asset.data.body_link_state_w[:, asset_cfg.body_ids[0], 3:7] # type: ignore return quat_error_magnitude(curr_quat_w, des_quat_w) diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/stack/mdp/franka_stack_events.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/stack/mdp/franka_stack_events.py index f741470a34..8d6aa87351 100644 --- a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/stack/mdp/franka_stack_events.py +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/stack/mdp/franka_stack_events.py @@ -130,10 +130,10 @@ def randomize_object_pose( pose_tensor = torch.tensor([pose_list[i]], device=env.device) positions = pose_tensor[:, 0:3] + env.scene.env_origins[cur_env, 0:3] orientations = math_utils.quat_from_euler_xyz(pose_tensor[:, 3], pose_tensor[:, 4], pose_tensor[:, 5]) - asset.write_root_pose_to_sim( + asset.write_root_link_pose_to_sim( torch.cat([positions, orientations], dim=-1), env_ids=torch.tensor([cur_env], device=env.device) ) - asset.write_root_velocity_to_sim( + asset.write_root_com_velocity_to_sim( torch.zeros(1, 6, device=env.device), env_ids=torch.tensor([cur_env], device=env.device) ) diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/stack/mdp/observations.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/stack/mdp/observations.py index 54831e9d75..53362934bb 100644 --- a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/stack/mdp/observations.py +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/stack/mdp/observations.py @@ -27,7 +27,7 @@ def cube_positions_in_world_frame( cube_2: RigidObject = env.scene[cube_2_cfg.name] cube_3: RigidObject = env.scene[cube_3_cfg.name] - return torch.cat((cube_1.data.root_pos_w, cube_2.data.root_pos_w, cube_3.data.root_pos_w), dim=1) + return torch.cat((cube_1.data.root_link_pos_w, cube_2.data.root_link_pos_w, cube_3.data.root_link_pos_w), dim=1) def instance_randomize_cube_positions_in_world_frame( @@ -69,7 +69,7 @@ def cube_orientations_in_world_frame( cube_2: RigidObject = env.scene[cube_2_cfg.name] cube_3: RigidObject = env.scene[cube_3_cfg.name] - return torch.cat((cube_1.data.root_quat_w, cube_2.data.root_quat_w, cube_3.data.root_quat_w), dim=1) + return torch.cat((cube_1.data.root_link_quat_w, cube_2.data.root_link_quat_w, cube_3.data.root_link_quat_w), dim=1) def instance_randomize_cube_orientations_in_world_frame( @@ -127,14 +127,14 @@ 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 = cube_1.data.root_pos_w - cube_1_quat_w = cube_1.data.root_quat_w + cube_1_pos_w = cube_1.data.root_link_pos_w + cube_1_quat_w = cube_1.data.root_link_quat_w - cube_2_pos_w = cube_2.data.root_pos_w - cube_2_quat_w = cube_2.data.root_quat_w + cube_2_pos_w = cube_2.data.root_link_pos_w + cube_2_quat_w = cube_2.data.root_link_quat_w - cube_3_pos_w = cube_3.data.root_pos_w - cube_3_quat_w = cube_3.data.root_quat_w + cube_3_pos_w = cube_3.data.root_link_pos_w + cube_3_quat_w = cube_3.data.root_link_quat_w ee_pos_w = ee_frame.data.target_pos_w[:, 0, :] gripper_to_cube_1 = cube_1_pos_w - ee_pos_w @@ -279,7 +279,7 @@ def object_grasped( ee_frame: FrameTransformer = env.scene[ee_frame_cfg.name] object: RigidObject = env.scene[object_cfg.name] - object_pos = object.data.root_pos_w + object_pos = object.data.root_link_pos_w end_effector_pos = ee_frame.data.target_pos_w[:, 0, :] pose_diff = torch.linalg.vector_norm(object_pos - end_effector_pos, dim=1) @@ -310,7 +310,7 @@ def object_stacked( upper_object: RigidObject = env.scene[upper_object_cfg.name] lower_object: RigidObject = env.scene[lower_object_cfg.name] - pos_diff = upper_object.data.root_pos_w - lower_object.data.root_pos_w + pos_diff = upper_object.data.root_link_pos_w - lower_object.data.root_link_pos_w height_dist = torch.linalg.vector_norm(pos_diff[:, 2:], dim=1) xy_dist = torch.linalg.vector_norm(pos_diff[:, :2], dim=1) diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/stack/mdp/terminations.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/stack/mdp/terminations.py index 9e5e8e161a..b4972006fb 100644 --- a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/stack/mdp/terminations.py +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/stack/mdp/terminations.py @@ -39,8 +39,8 @@ def cubes_stacked( cube_2: RigidObject = env.scene[cube_2_cfg.name] cube_3: RigidObject = env.scene[cube_3_cfg.name] - pos_diff_c12 = cube_1.data.root_pos_w - cube_2.data.root_pos_w - pos_diff_c23 = cube_2.data.root_pos_w - cube_3.data.root_pos_w + pos_diff_c12 = cube_1.data.root_link_pos_w - cube_2.data.root_link_pos_w + pos_diff_c23 = cube_2.data.root_link_pos_w - cube_3.data.root_link_pos_w # Compute cube position difference in x-y plane xy_dist_c12 = torch.norm(pos_diff_c12[:, :2], dim=1) diff --git a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/navigation/mdp/pre_trained_policy_action.py b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/navigation/mdp/pre_trained_policy_action.py index 7e5af7d313..1936ffa61a 100644 --- a/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/navigation/mdp/pre_trained_policy_action.py +++ b/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/navigation/mdp/pre_trained_policy_action.py @@ -128,11 +128,11 @@ def _debug_vis_callback(self, event): return # get marker location # -- base state - base_pos_w = self.robot.data.root_pos_w.clone() + base_pos_w = self.robot.data.root_link_pos_w.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(self.robot.data.root_lin_vel_b[:, :2]) + vel_arrow_scale, vel_arrow_quat = self._resolve_xy_velocity_to_arrow(self.robot.data.root_com_lin_vel_b[:, :2]) # display markers self.base_vel_goal_visualizer.visualize(base_pos_w, vel_des_arrow_quat, vel_des_arrow_scale) self.base_vel_visualizer.visualize(base_pos_w, vel_arrow_quat, vel_arrow_scale) @@ -153,7 +153,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 = self.robot.data.root_quat_w + base_quat_w = self.robot.data.root_link_quat_w arrow_quat = math_utils.quat_mul(base_quat_w, arrow_quat) return arrow_scale, arrow_quat diff --git a/source/standalone/benchmarks/benchmark_load_robot.py b/source/standalone/benchmarks/benchmark_load_robot.py index 3ac0345e37..0704397244 100644 --- a/source/standalone/benchmarks/benchmark_load_robot.py +++ b/source/standalone/benchmarks/benchmark_load_robot.py @@ -115,7 +115,8 @@ def run_simulator(sim: sim_utils.SimulationContext, scene: InteractiveScene): # if this is not done, then the robots will be spawned at the (0, 0, 0) of the simulation world root_state = robot.data.default_root_state.clone() root_state[:, :3] += scene.env_origins - robot.write_root_state_to_sim(root_state) + robot.write_root_link_pose_to_sim(root_state[:, :7]) + robot.write_root_com_velocity_to_sim(root_state[:, 7:]) # set joint positions with some noise joint_pos, joint_vel = robot.data.default_joint_pos.clone(), robot.data.default_joint_vel.clone() joint_pos += torch.rand_like(joint_pos) * 0.1 diff --git a/source/standalone/demos/arms.py b/source/standalone/demos/arms.py index d50034d4d9..d68cde9f6c 100644 --- a/source/standalone/demos/arms.py +++ b/source/standalone/demos/arms.py @@ -179,7 +179,8 @@ def run_simulator(sim: sim_utils.SimulationContext, entities: dict[str, Articula # root state root_state = robot.data.default_root_state.clone() root_state[:, :3] += origins[index] - robot.write_root_state_to_sim(root_state) + robot.write_root_link_pose_to_sim(root_state[:, :7]) + robot.write_root_com_velocity_to_sim(root_state[:, 7:]) # set joint positions joint_pos, joint_vel = robot.data.default_joint_pos.clone(), robot.data.default_joint_vel.clone() robot.write_joint_state_to_sim(joint_pos, joint_vel) diff --git a/source/standalone/demos/bipeds.py b/source/standalone/demos/bipeds.py index 961af39e19..0ec291be74 100644 --- a/source/standalone/demos/bipeds.py +++ b/source/standalone/demos/bipeds.py @@ -97,7 +97,8 @@ def main(): robot.write_joint_state_to_sim(joint_pos, joint_vel) root_state = robot.data.default_root_state.clone() root_state[:, :3] += origins[index] - robot.write_root_state_to_sim(root_state) + robot.write_root_link_pose_to_sim(root_state[:, :7]) + robot.write_root_com_velocity_to_sim(root_state[:, 7:]) robot.reset() # reset command print(">>>>>>>> Reset!") diff --git a/source/standalone/demos/cameras.py b/source/standalone/demos/cameras.py index 6c5583d55a..204d8436df 100644 --- a/source/standalone/demos/cameras.py +++ b/source/standalone/demos/cameras.py @@ -187,7 +187,8 @@ def run_simulator(sim: sim_utils.SimulationContext, scene: InteractiveScene): # if this is not done, then the robots will be spawned at the (0, 0, 0) of the simulation world root_state = scene["robot"].data.default_root_state.clone() root_state[:, :3] += scene.env_origins - scene["robot"].write_root_state_to_sim(root_state) + scene["robot"].write_root_link_pose_to_sim(root_state[:, :7]) + scene["robot"].write_root_com_velocity_to_sim(root_state[:, 7:]) # set joint positions with some noise joint_pos, joint_vel = ( scene["robot"].data.default_joint_pos.clone(), diff --git a/source/standalone/demos/hands.py b/source/standalone/demos/hands.py index e939d63cf5..9d67c3665e 100644 --- a/source/standalone/demos/hands.py +++ b/source/standalone/demos/hands.py @@ -113,7 +113,8 @@ def run_simulator(sim: sim_utils.SimulationContext, entities: dict[str, Articula # root state root_state = robot.data.default_root_state.clone() root_state[:, :3] += origins[index] - robot.write_root_state_to_sim(root_state) + robot.write_root_link_pose_to_sim(root_state[:, :7]) + robot.write_root_com_velocity_to_sim(root_state[:, 7:]) # joint state joint_pos, joint_vel = robot.data.default_joint_pos.clone(), robot.data.default_joint_vel.clone() robot.write_joint_state_to_sim(joint_pos, joint_vel) diff --git a/source/standalone/demos/multi_asset.py b/source/standalone/demos/multi_asset.py index 2df770d13f..e640e976d2 100644 --- a/source/standalone/demos/multi_asset.py +++ b/source/standalone/demos/multi_asset.py @@ -240,16 +240,19 @@ def run_simulator(sim: SimulationContext, scene: InteractiveScene): # object root_state = rigid_object.data.default_root_state.clone() root_state[:, :3] += scene.env_origins - rigid_object.write_root_state_to_sim(root_state) + rigid_object.write_root_link_pose_to_sim(root_state[:, :7]) + rigid_object.write_root_com_velocity_to_sim(root_state[:, 7:]) # object collection object_state = rigid_object_collection.data.default_object_state.clone() object_state[..., :3] += scene.env_origins.unsqueeze(1) - rigid_object_collection.write_object_state_to_sim(object_state) + rigid_object_collection.write_object_link_pose_to_sim(object_state[..., :7]) + rigid_object_collection.write_object_com_velocity_to_sim(object_state[..., 7:]) # robot # -- root state root_state = robot.data.default_root_state.clone() root_state[:, :3] += scene.env_origins - robot.write_root_state_to_sim(root_state) + robot.write_root_link_pose_to_sim(root_state[:, :7]) + robot.write_root_com_velocity_to_sim(root_state[:, 7:]) # -- joint state joint_pos, joint_vel = robot.data.default_joint_pos.clone(), robot.data.default_joint_vel.clone() robot.write_joint_state_to_sim(joint_pos, joint_vel) diff --git a/source/standalone/demos/quadcopter.py b/source/standalone/demos/quadcopter.py index 2110ae1219..88e00745fc 100644 --- a/source/standalone/demos/quadcopter.py +++ b/source/standalone/demos/quadcopter.py @@ -91,8 +91,8 @@ def main(): # reset dof state joint_pos, joint_vel = robot.data.default_joint_pos, robot.data.default_joint_vel robot.write_joint_state_to_sim(joint_pos, joint_vel) - robot.write_root_pose_to_sim(robot.data.default_root_state[:, :7]) - robot.write_root_velocity_to_sim(robot.data.default_root_state[:, 7:]) + robot.write_root_link_pose_to_sim(robot.data.default_root_state[:, :7]) + robot.write_root_com_velocity_to_sim(robot.data.default_root_state[:, 7:]) robot.reset() # reset command print(">>>>>>>> Reset!") diff --git a/source/standalone/demos/quadrupeds.py b/source/standalone/demos/quadrupeds.py index 3122a4741c..8cabf4e92b 100644 --- a/source/standalone/demos/quadrupeds.py +++ b/source/standalone/demos/quadrupeds.py @@ -142,7 +142,8 @@ def run_simulator(sim: sim_utils.SimulationContext, entities: dict[str, Articula # root state root_state = robot.data.default_root_state.clone() root_state[:, :3] += origins[index] - robot.write_root_state_to_sim(root_state) + robot.write_root_link_pose_to_sim(root_state[:, :7]) + robot.write_root_com_velocity_to_sim(root_state[:, 7:]) # joint state joint_pos, joint_vel = robot.data.default_joint_pos.clone(), robot.data.default_joint_vel.clone() robot.write_joint_state_to_sim(joint_pos, joint_vel) diff --git a/source/standalone/demos/sensors/contact_sensor.py b/source/standalone/demos/sensors/contact_sensor.py index 35de7fe667..1559bdd701 100644 --- a/source/standalone/demos/sensors/contact_sensor.py +++ b/source/standalone/demos/sensors/contact_sensor.py @@ -109,7 +109,8 @@ def run_simulator(sim: sim_utils.SimulationContext, scene: InteractiveScene): # if this is not done, then the robots will be spawned at the (0, 0, 0) of the simulation world root_state = scene["robot"].data.default_root_state.clone() root_state[:, :3] += scene.env_origins - scene["robot"].write_root_state_to_sim(root_state) + scene["robot"].write_root_link_pose_to_sim(root_state[:, :7]) + scene["robot"].write_root_com_velocity_to_sim(root_state[:, 7:]) # set joint positions with some noise joint_pos, joint_vel = ( scene["robot"].data.default_joint_pos.clone(), diff --git a/source/standalone/demos/sensors/frame_transformer_sensor.py b/source/standalone/demos/sensors/frame_transformer_sensor.py index e66e59f500..6d009a59d0 100644 --- a/source/standalone/demos/sensors/frame_transformer_sensor.py +++ b/source/standalone/demos/sensors/frame_transformer_sensor.py @@ -105,7 +105,8 @@ def run_simulator(sim: sim_utils.SimulationContext, scene: InteractiveScene): # if this is not done, then the robots will be spawned at the (0, 0, 0) of the simulation world root_state = scene["robot"].data.default_root_state.clone() root_state[:, :3] += scene.env_origins - scene["robot"].write_root_state_to_sim(root_state) + scene["robot"].write_root_link_pose_to_sim(root_state[:, :7]) + scene["robot"].write_root_com_velocity_to_sim(root_state[:, 7:]) # set joint positions with some noise joint_pos, joint_vel = ( scene["robot"].data.default_joint_pos.clone(), diff --git a/source/standalone/demos/sensors/raycaster_sensor.py b/source/standalone/demos/sensors/raycaster_sensor.py index db7d74d3de..846ee8a41c 100644 --- a/source/standalone/demos/sensors/raycaster_sensor.py +++ b/source/standalone/demos/sensors/raycaster_sensor.py @@ -63,6 +63,7 @@ class RaycasterSensorSceneCfg(InteractiveSceneCfg): update_period=1 / 60, offset=RayCasterCfg.OffsetCfg(pos=(0, 0, 0.5)), mesh_prim_paths=["/World/Ground"], + attach_yaw_only=True, pattern_cfg=patterns.LidarPatternCfg( channels=100, vertical_fov_range=[-90, 90], horizontal_fov_range=[-90, 90], horizontal_res=1.0 ), @@ -92,7 +93,8 @@ def run_simulator(sim: sim_utils.SimulationContext, scene: InteractiveScene): # if this is not done, then the robots will be spawned at the (0, 0, 0) of the simulation world root_state = scene["robot"].data.default_root_state.clone() root_state[:, :3] += scene.env_origins - scene["robot"].write_root_state_to_sim(root_state) + scene["robot"].write_root_link_pose_to_sim(root_state[:, :7]) + scene["robot"].write_root_com_velocity_to_sim(root_state[:, 7:]) # set joint positions with some noise joint_pos, joint_vel = ( scene["robot"].data.default_joint_pos.clone(), diff --git a/source/standalone/environments/state_machine/lift_cube_sm.py b/source/standalone/environments/state_machine/lift_cube_sm.py index 03f9c16e5f..503f5d8c0c 100644 --- a/source/standalone/environments/state_machine/lift_cube_sm.py +++ b/source/standalone/environments/state_machine/lift_cube_sm.py @@ -293,7 +293,7 @@ def main(): tcp_rest_orientation = ee_frame_sensor.data.target_quat_w[..., 0, :].clone() # -- object frame object_data: RigidObjectData = env.unwrapped.scene["object"].data - object_position = object_data.root_pos_w - env.unwrapped.scene.env_origins + object_position = object_data.root_link_pos_w - env.unwrapped.scene.env_origins # -- target object frame desired_position = env.unwrapped.command_manager.get_command("object_pose")[..., :3] diff --git a/source/standalone/tutorials/01_assets/run_articulation.py b/source/standalone/tutorials/01_assets/run_articulation.py index d06874ac40..bcfa409046 100644 --- a/source/standalone/tutorials/01_assets/run_articulation.py +++ b/source/standalone/tutorials/01_assets/run_articulation.py @@ -94,7 +94,8 @@ def run_simulator(sim: sim_utils.SimulationContext, entities: dict[str, Articula # if this is not done, then the robots will be spawned at the (0, 0, 0) of the simulation world root_state = robot.data.default_root_state.clone() root_state[:, :3] += origins - robot.write_root_state_to_sim(root_state) + robot.write_root_link_pose_to_sim(root_state[:, :7]) + robot.write_root_com_velocity_to_sim(root_state[:, 7:]) # set joint positions with some noise joint_pos, joint_vel = robot.data.default_joint_pos.clone(), robot.data.default_joint_vel.clone() joint_pos += torch.rand_like(joint_pos) * 0.1 diff --git a/source/standalone/tutorials/01_assets/run_rigid_object.py b/source/standalone/tutorials/01_assets/run_rigid_object.py index cede00b2c1..2f5cd327ce 100644 --- a/source/standalone/tutorials/01_assets/run_rigid_object.py +++ b/source/standalone/tutorials/01_assets/run_rigid_object.py @@ -103,7 +103,8 @@ def run_simulator(sim: sim_utils.SimulationContext, entities: dict[str, RigidObj radius=0.1, h_range=(0.25, 0.5), size=cone_object.num_instances, device=cone_object.device ) # write root state to simulation - cone_object.write_root_state_to_sim(root_state) + cone_object.write_root_link_pose_to_sim(root_state[:, :7]) + cone_object.write_root_com_velocity_to_sim(root_state[:, 7:]) # reset buffers cone_object.reset() print("----------------------------------------") @@ -119,7 +120,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_link_state_w[:, :3]}") def main(): diff --git a/source/standalone/tutorials/02_scene/create_scene.py b/source/standalone/tutorials/02_scene/create_scene.py index 36096d35d5..549d820508 100644 --- a/source/standalone/tutorials/02_scene/create_scene.py +++ b/source/standalone/tutorials/02_scene/create_scene.py @@ -83,7 +83,8 @@ def run_simulator(sim: sim_utils.SimulationContext, scene: InteractiveScene): # if this is not done, then the robots will be spawned at the (0, 0, 0) of the simulation world root_state = robot.data.default_root_state.clone() root_state[:, :3] += scene.env_origins - robot.write_root_state_to_sim(root_state) + robot.write_root_link_pose_to_sim(root_state[:, :7]) + robot.write_root_com_velocity_to_sim(root_state[:, 7:]) # set joint positions with some noise joint_pos, joint_vel = robot.data.default_joint_pos.clone(), robot.data.default_joint_vel.clone() joint_pos += torch.rand_like(joint_pos) * 0.1 diff --git a/source/standalone/tutorials/03_envs/create_cube_base_env.py b/source/standalone/tutorials/03_envs/create_cube_base_env.py index 92a47ed6cf..ee6643c396 100644 --- a/source/standalone/tutorials/03_envs/create_cube_base_env.py +++ b/source/standalone/tutorials/03_envs/create_cube_base_env.py @@ -120,11 +120,11 @@ 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 - (self._asset.data.root_pos_w - self._env.scene.env_origins) - vel_error = -self._asset.data.root_lin_vel_w + pos_error = self._processed_actions - (self._asset.data.root_link_pos_w - self._env.scene.env_origins) + vel_error = -self._asset.data.root_com_lin_vel_w # 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) + self._asset.write_root_com_velocity_to_sim(self._vel_command) @configclass @@ -149,7 +149,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 asset.data.root_pos_w - env.scene.env_origins + return asset.data.root_link_pos_w - env.scene.env_origins ## diff --git a/source/standalone/tutorials/04_sensors/add_sensors_on_robot.py b/source/standalone/tutorials/04_sensors/add_sensors_on_robot.py index 8ba29c2921..8c367a1230 100644 --- a/source/standalone/tutorials/04_sensors/add_sensors_on_robot.py +++ b/source/standalone/tutorials/04_sensors/add_sensors_on_robot.py @@ -113,7 +113,8 @@ def run_simulator(sim: sim_utils.SimulationContext, scene: InteractiveScene): # if this is not done, then the robots will be spawned at the (0, 0, 0) of the simulation world root_state = scene["robot"].data.default_root_state.clone() root_state[:, :3] += scene.env_origins - scene["robot"].write_root_state_to_sim(root_state) + scene["robot"].write_root_link_pose_to_sim(root_state[:, :7]) + scene["robot"].write_root_com_velocity_to_sim(root_state[:, 7:]) # set joint positions with some noise joint_pos, joint_vel = ( scene["robot"].data.default_joint_pos.clone(), diff --git a/source/standalone/tutorials/04_sensors/run_ray_caster.py b/source/standalone/tutorials/04_sensors/run_ray_caster.py index 921eebad53..c56c2f9ec7 100644 --- a/source/standalone/tutorials/04_sensors/run_ray_caster.py +++ b/source/standalone/tutorials/04_sensors/run_ray_caster.py @@ -109,7 +109,8 @@ def run_simulator(sim: sim_utils.SimulationContext, scene_entities: dict): # Reset the scene if step_count % 250 == 0: # reset the balls - balls.write_root_state_to_sim(ball_default_state) + balls.write_root_link_pose_to_sim(ball_default_state[:, :7]) + balls.write_root_com_velocity_to_sim(ball_default_state[:, 7:]) # reset the sensor ray_caster.reset() # reset the counter diff --git a/source/standalone/tutorials/05_controllers/run_diff_ik.py b/source/standalone/tutorials/05_controllers/run_diff_ik.py index 51c2102070..68483f2902 100644 --- a/source/standalone/tutorials/05_controllers/run_diff_ik.py +++ b/source/standalone/tutorials/05_controllers/run_diff_ik.py @@ -160,8 +160,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_link_state_w[:, robot_entity_cfg.body_ids[0], 0:7] + root_pose_w = robot.data.root_link_state_w[:, 0:7] 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( @@ -181,7 +181,7 @@ def run_simulator(sim: sim_utils.SimulationContext, scene: InteractiveScene): scene.update(sim_dt) # obtain quantities from simulation - ee_pose_w = robot.data.body_state_w[:, robot_entity_cfg.body_ids[0], 0:7] + ee_pose_w = robot.data.body_link_state_w[:, 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/source/standalone/tutorials/05_controllers/run_osc.py b/source/standalone/tutorials/05_controllers/run_osc.py index 8970d364a2..252b93ce01 100644 --- a/source/standalone/tutorials/05_controllers/run_osc.py +++ b/source/standalone/tutorials/05_controllers/run_osc.py @@ -291,24 +291,26 @@ def update_states( gravity = robot.root_physx_view.get_generalized_gravity_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(robot.data.root_state_w[:, 3:7])) + root_rot_matrix = matrix_from_quat(quat_inv(robot.data.root_link_quat_w)) 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 = robot.data.root_state_w[:, 0:7] - ee_pose_w = robot.data.body_state_w[:, ee_frame_idx, 0:7] - 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] - ) + root_pos_w = robot.data.root_link_pos_w + root_quat_w = robot.data.root_link_quat_w + ee_pos_w = robot.data.body_link_pos_w[:, ee_frame_idx] + ee_quat_w = robot.data.body_link_quat_w[:, 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 = robot.data.body_vel_w[:, ee_frame_idx, :] # Extract end-effector velocity in the world frame - root_vel_w = robot.data.root_vel_w # Extract root velocity in the world frame + ee_vel_w = robot.data.body_com_vel_w[:, ee_frame_idx, :] # Extract end-effector velocity in the world frame + root_vel_w = robot.data.root_com_vel_w # 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_rotate_inverse(robot.data.root_quat_w, relative_vel_w[:, 0:3]) # From world to root frame - ee_ang_vel_b = quat_rotate_inverse(robot.data.root_quat_w, relative_vel_w[:, 3:6]) + ee_lin_vel_b = quat_rotate_inverse(robot.data.root_link_quat_w, relative_vel_w[:, 0:3]) # From world to root frame + ee_ang_vel_b = quat_rotate_inverse(robot.data.root_link_quat_w, relative_vel_w[:, 3:6]) ee_vel_b = torch.cat([ee_lin_vel_b, ee_ang_vel_b], dim=-1) # Calculate the contact force