Unverified Commit 999c1e9a authored by James Tigue's avatar James Tigue Committed by GitHub

Removes deprecation for root_state_w properties and setters (#1695)

…ection root and base states

# Description

<!--
Thank you for your interest in sending a pull request. Please make sure
to check the contribution guidelines.

Link:
https://isaac-sim.github.io/IsaacLab/main/source/refs/contributing.html
-->

This PR removes the deprecation of root_state_w setters and properties
due to significant regression in training speed. It leaves it the
root_link_* and root_com_* properties and setters but does not force
them on examples and tests.

Fixes #1699 

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

## Type of change

<!-- As you go through the list, delete the ones that are not
applicable. -->

- Bug fix (non-breaking change which fixes an issue)

<!--
Example:

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

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

## Checklist

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

<!--
As you go through the checklist above, you can mark something as done by
putting an x character in it

For example,
- [x] I have done this task
- [ ] I have not done this task
-->

---------
Signed-off-by: 's avatarJames Tigue <166445701+jtigue-bdai@users.noreply.github.com>
Co-authored-by: 's avatarKelly Guo <kellyg@nvidia.com>
Co-authored-by: 's avatarKelly Guo <kellyguo123@hotmail.com>
parent 7ed94ce5
......@@ -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_link_pose_to_sim( |
| | self.cartpole.write_root_pose_to_sim( |
| | default_root_state[:, :7], env_ids) |
| | self.cartpole.write_root_com_velocity_to_sim( |
| | self.cartpole.write_root_velocity_to_sim( |
| | default_root_state[:, 7:], env_ids) |
| | self.cartpole.write_joint_state_to_sim( |
| | joint_pos, joint_vel, None, env_ids) |
......
......@@ -364,8 +364,8 @@ In Isaac Lab, ``root_pose`` and ``root_velocity`` have been combined into single
.. code-block::python
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_root_pose_to_sim(default_root_state[:, :7], env_ids)
self.cartpole.write_root_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_link_pose_to_sim( |
| indices = env_ids.to(dtype=torch.int32) | self.cartpole.write_root_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_com_velocity_to_sim( |
| self._cartpoles.set_joint_velocities(dof_vel, indices=indices) | self.cartpole.write_root_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) |
......
......@@ -66,7 +66,7 @@ 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_link_pose_to_sim` and :meth:`Articulation.write_root_com_velocity_to_sim`
To reset the articulation, we first set the root state by calling the :meth:`Articulation.write_root_pose_to_sim` and :meth:`Articulation.write_root_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.
......
......@@ -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_link_pos_w[:, :3]}")
:end-at: print(f"Root position (in world): {cube_object.data.root_pos_w[:, :3]}")
The Code Execution
......
......@@ -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_link_pose_to_sim` and :meth:`assets.RigidObject.write_root_com_velocity_to_sim` methods.
set the desired state of the rigid object prim using the :meth:`assets.RigidObject.write_root_pose_to_sim` and :meth:`assets.RigidObject.write_root_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
......
......@@ -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_link_state_w` provides the state of all the
the attribute :attr:`assets.Articulationdata.body_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.
......
[package]
# Note: Semantic Versioning is used: https://semver.org/
version = "0.30.5"
version = "0.30.6"
# Description
title = "Isaac Lab framework for Robot Learning"
......
Changelog
---------
0.30.6 (2025-01-17)
~~~~~~~~~~~~~~~~~~~
Fixed
^^^^^
* removed deprecation of :attr:`omni.isaac.lab.assets.ArticulationData.root_state_w` and
:attr:`omni.isaac.lab.assets.ArticulationData.body_state_w` derived properties.
* removed deprecation of :meth:`omni.isaac.lab.assets.Articulation.write_root_state_to_sim`.
* replaced calls to :attr:`omni.isaac.lab.assets.ArticulationData.root_com_state_w` and
:attr:`omni.isaac.lab.assets.ArticulationData.root_link_state_w` with corresponding calls to
:attr:`omni.isaac.lab.assets.ArticulationData.root_state_w`.
* replaced calls to :attr:`omni.isaac.lab.assets.ArticulationData.body_com_state_w` and
:attr:`omni.isaac.lab.assets.ArticulationData.body_link_state_w` properties with corresponding calls to
:attr:`omni.isaac.lab.assets.ArticulationData.body_state_w` properties.
* removed deprecation of :attr:`omni.isaac.lab.assets.RigidObjectData.root_state_w` derived properties .
* removed deprecation of :meth:`omni.isaac.lab.assets.RigidObject.write_root_state_to_sim`.
* replaced calls to :attr:`omni.isaac.lab.assets.RigidObjectData.root_com_state_w` and
:attr:`omni.isaac.lab.assets.RigidObjectData.root_link_state_w` properties with corresponding calls to
:attr:`omni.isaac.lab.assets.RigidObjectData.root_state_w` properties.
* removed deprecation of :attr:`omni.isaac.lab.assets.RigidObjectCollectionData.root_state_w` derived properties.
* removed deprecation of :meth:`omni.isaac.lab.assets.RigidObjectCollection.write_root_state_to_sim`.
* replaced calls to :attr:`omni.isaac.lab.assets.RigidObjectCollectionData.root_com_state_w` and
:attr:`omni.isaac.lab.assets.RigidObjectData.root_link_state_w` properties with corresponding calls to
:attr:`omni.isaac.lab.assets.RigidObjectData.root_state_w` properties.
* fixed indexing issue in ``write_root_link_velocity_to_sim`` in :class:`omni.isaac.lab.assets.RigidObject`
* fixed index broadcasting in ``write_object_link_velocity_to_sim`` and ``write_object_com_pose_to_sim`` in :class:`omni.isaac.lab.assets.RigidObjectCollection`
0.30.5 (2025-01-14)
~~~~~~~~~~~~~~~~~~~
......
......@@ -98,10 +98,6 @@ class Articulation(AssetBase):
"""
super().__init__(cfg)
self._root_state_dep_warn = False
self._root_pose_dep_warn = False
self._root_vel_dep_warn = False
"""
Properties
"""
......@@ -285,14 +281,6 @@ class Articulation(AssetBase):
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)
......@@ -334,15 +322,21 @@ class Articulation(AssetBase):
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)
# resolve all indices
physx_env_ids = env_ids
if env_ids is None:
env_ids = slice(None)
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()
# convert root quaternion from wxyz to xyzw
root_poses_xyzw = self._data.root_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
# set into simulation
self.root_physx_view.set_root_transforms(root_poses_xyzw, indices=physx_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.
......@@ -361,9 +355,7 @@ class Articulation(AssetBase):
# note: we need to do this here since tensors are not set into simulation until step.
# set into internal buffers
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_link_state_w[:, :7].clone()
root_poses_xyzw[:, 3:] = math_utils.convert_quat(root_poses_xyzw[:, 3:], to="xyzw")
......@@ -413,15 +405,17 @@ class Articulation(AssetBase):
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)
# resolve all indices
physx_env_ids = env_ids
if env_ids is None:
env_ids = slice(None)
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.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)
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.
......@@ -442,9 +436,7 @@ class Articulation(AssetBase):
# note: we need to do this here since tensors are not set into simulation until step.
# set into internal buffers
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_com_state_w[:, 7:], indices=physx_env_ids)
......@@ -507,8 +499,8 @@ class Articulation(AssetBase):
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
# 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)
......
......@@ -6,7 +6,6 @@
import torch
import weakref
import omni.log
import omni.physics.tensors.impl.api as physx
import omni.isaac.lab.utils.math as math_utils
......@@ -75,11 +74,6 @@ class ArticulationData:
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
......@@ -274,13 +268,6 @@ class ArticulationData:
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()
......@@ -347,13 +334,6 @@ class ArticulationData:
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
......
......@@ -55,9 +55,6 @@ class RigidObject(AssetBase):
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
......@@ -159,13 +156,6 @@ class RigidObject(AssetBase):
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)
......@@ -208,15 +198,19 @@ class RigidObject(AssetBase):
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)
# resolve all indices
physx_env_ids = env_ids
if env_ids is None:
env_ids = slice(None)
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()
# convert root quaternion from wxyz to xyzw
root_poses_xyzw = self._data.root_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_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.
......@@ -235,9 +229,7 @@ class RigidObject(AssetBase):
# note: we need to do this here since tensors are not set into simulation until step.
# set into internal buffers
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_link_state_w[:, :7].clone()
root_poses_xyzw[:, 3:] = math_utils.convert_quat(root_poses_xyzw[:, 3:], to="xyzw")
......@@ -254,10 +246,11 @@ class RigidObject(AssetBase):
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)
local_env_ids = slice(env_ids)
else:
local_env_ids = env_ids
com_pos = self.data.com_pos_b[local_env_ids, 0, :]
com_quat = self.data.com_quat_b[local_env_ids, 0, :]
......@@ -282,15 +275,17 @@ class RigidObject(AssetBase):
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)
# resolve all indices
physx_env_ids = env_ids
if env_ids is None:
env_ids = slice(None)
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.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)
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.
......@@ -311,9 +306,7 @@ class RigidObject(AssetBase):
# note: we need to do this here since tensors are not set into simulation until step.
# set into internal buffers
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_com_state_w[:, 7:], indices=physx_env_ids)
......@@ -330,7 +323,9 @@ class RigidObject(AssetBase):
"""
# resolve all indices
if env_ids is None:
local_env_ids = slice(None)
local_env_ids = slice(env_ids)
else:
local_env_ids = env_ids
root_com_velocity = root_velocity.clone()
quat = self.data.root_link_state_w[local_env_ids, 3:7]
......
......@@ -69,10 +69,6 @@ class RigidObjectData:
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.
......@@ -121,12 +117,6 @@ class RigidObjectData:
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
......
......@@ -105,10 +105,6 @@ class RigidObjectCollection(AssetBase):
self._debug_vis_handle = None
self._root_state_dep_warn = False
self._root_pose_dep_warn = False
self._root_vel_dep_warn = False
"""
Properties
"""
......@@ -226,13 +222,6 @@ class RigidObjectCollection(AssetBase):
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)
......@@ -291,15 +280,22 @@ class RigidObjectCollection(AssetBase):
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)
# resolve all indices
# -- env_ids
if env_ids is None:
env_ids = self._ALL_ENV_INDICES
# -- object_ids
if object_ids is None:
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_state_w[env_ids[:, None], object_ids, :7] = object_pose.clone()
# convert the quaternion from wxyz to xyzw
poses_xyzw = self._data.object_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_link_pose_to_sim(
self,
......@@ -327,9 +323,7 @@ class RigidObjectCollection(AssetBase):
# 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_link_state_w[..., :7].clone()
poses_xyzw[..., 3:] = math_utils.convert_quat(poses_xyzw[..., 3:], to="xyzw")
......@@ -362,8 +356,8 @@ class RigidObjectCollection(AssetBase):
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, :]
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],
......@@ -388,15 +382,22 @@ class RigidObjectCollection(AssetBase):
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
# resolve all indices
# -- env_ids
if env_ids is None:
env_ids = self._ALL_ENV_INDICES
# -- object_ids
if object_ids is None:
object_ids = self._ALL_OBJ_INDICES
self._data.object_state_w[env_ids[:, None], object_ids, 7:] = object_velocity.clone()
self._data.object_acc_w[env_ids[:, None], object_ids] = 0.0
self.write_object_com_velocity_to_sim(object_velocity=object_velocity, env_ids=env_ids, object_ids=object_ids)
# 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
)
def write_object_com_velocity_to_sim(
self,
......@@ -420,9 +421,7 @@ class RigidObjectCollection(AssetBase):
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
......@@ -456,8 +455,8 @@ class RigidObjectCollection(AssetBase):
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, :]
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
......
......@@ -6,7 +6,6 @@
import torch
import weakref
import omni.log
import omni.physics.tensors.impl.api as physx
import omni.isaac.lab.utils.math as math_utils
......@@ -74,10 +73,6 @@ class RigidObjectCollectionData:
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.
......@@ -127,13 +122,6 @@ class RigidObjectCollectionData:
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())
......
......@@ -150,7 +150,7 @@ class NonHolonomicAction(ActionTerm):
def apply_actions(self):
# obtain current heading
quat_w = self._asset.data.body_link_quat_w[:, self._body_idx].view(self.num_envs, 4)
quat_w = self._asset.data.body_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
......
......@@ -142,7 +142,7 @@ class DifferentialInverseKinematicsAction(ActionTerm):
@property
def jacobian_b(self) -> torch.Tensor:
jacobian = self.jacobian_w
base_rot = self._asset.data.root_link_quat_w
base_rot = self._asset.data.root_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,10 +192,10 @@ class DifferentialInverseKinematicsAction(ActionTerm):
A tuple of the body's position and orientation in the root frame.
"""
# obtain quantities from simulation
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
ee_pos_w = self._asset.data.body_pos_w[:, self._body_idx]
ee_quat_w = self._asset.data.body_quat_w[:, self._body_idx]
root_pos_w = self._asset.data.root_pos_w
root_quat_w = self._asset.data.root_quat_w
# compute the pose of the body in the root frame
ee_pose_b, ee_quat_b = math_utils.subtract_frame_transforms(root_pos_w, root_quat_w, ee_pos_w, ee_quat_w)
# account for the offset
......@@ -403,7 +403,7 @@ class OperationalSpaceControllerAction(ActionTerm):
@property
def jacobian_b(self) -> torch.Tensor:
jacobian = self.jacobian_w
base_rot = self._asset.data.root_link_quat_w
base_rot = self._asset.data.root_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:, :])
......@@ -597,12 +597,12 @@ class OperationalSpaceControllerAction(ActionTerm):
def _compute_ee_pose(self):
"""Computes the pose of the ee frame in root frame."""
# Obtain quantities from simulation
self._ee_pose_w[:, 0:3] = 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]
self._ee_pose_w[:, 0:3] = self._asset.data.body_pos_w[:, self._ee_body_idx]
self._ee_pose_w[:, 3:7] = self._asset.data.body_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_link_pos_w,
self._asset.data.root_link_quat_w,
self._asset.data.root_pos_w,
self._asset.data.root_quat_w,
self._ee_pose_w[:, 0:3],
self._ee_pose_w[:, 3:7],
)
......@@ -617,17 +617,13 @@ class OperationalSpaceControllerAction(ActionTerm):
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_com_vel_w[:, self._ee_body_idx, :]
self._ee_vel_w[:] = self._asset.data.body_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_com_vel_w
relative_vel_w = self._ee_vel_w - self._asset.data.root_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_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]
)
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])
# Account for the offset
if self.cfg.body_offset is not None:
......@@ -644,7 +640,7 @@ class OperationalSpaceControllerAction(ActionTerm):
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_link_quat_w, self._ee_force_w)
self._ee_force_b[:] = math_utils.quat_rotate_inverse(self._asset.data.root_quat_w, self._ee_force_w)
def _compute_joint_states(self):
"""Computes the joint states for operational space control."""
......@@ -659,8 +655,8 @@ class OperationalSpaceControllerAction(ActionTerm):
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_link_pos_w,
self._asset.data.root_link_quat_w,
self._asset.data.root_pos_w,
self._asset.data.root_quat_w,
self._task_frame_transformer.data.target_pos_w[:, 0, :],
self._task_frame_transformer.data.target_quat_w[:, 0, :],
)
......
......@@ -81,9 +81,7 @@ class UniformPose2dCommand(CommandTerm):
def _update_metrics(self):
# logs data
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_pos_2d"] = torch.norm(self.pos_command_w[:, :2] - self.robot.data.root_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]):
......@@ -97,7 +95,7 @@ class UniformPose2dCommand(CommandTerm):
if self.cfg.simple_heading:
# set heading command to point towards target
target_vec = self.pos_command_w[env_ids] - self.robot.data.root_link_pos_w[env_ids]
target_vec = self.pos_command_w[env_ids] - self.robot.data.root_pos_w[env_ids]
target_direction = torch.atan2(target_vec[:, 1], target_vec[:, 0])
flipped_target_direction = wrap_to_pi(target_direction + torch.pi)
......@@ -118,8 +116,8 @@ class UniformPose2dCommand(CommandTerm):
def _update_command(self):
"""Re-target the position command to the current root state."""
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)
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)
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):
......@@ -184,7 +182,7 @@ class TerrainBasedPose2dCommand(UniformPose2dCommand):
if self.cfg.simple_heading:
# set heading command to point towards target
target_vec = self.pos_command_w[env_ids] - self.robot.data.root_link_pos_w[env_ids]
target_vec = self.pos_command_w[env_ids] - self.robot.data.root_pos_w[env_ids]
target_direction = torch.atan2(target_vec[:, 1], target_vec[:, 0])
flipped_target_direction = wrap_to_pi(target_direction + torch.pi)
......
......@@ -92,8 +92,8 @@ class UniformPoseCommand(CommandTerm):
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_link_pos_w,
self.robot.data.root_link_quat_w,
self.robot.data.root_pos_w,
self.robot.data.root_quat_w,
self.pose_command_b[:, :3],
self.pose_command_b[:, 3:],
)
......@@ -101,8 +101,8 @@ class UniformPoseCommand(CommandTerm):
pos_error, rot_error = compute_pose_error(
self.pose_command_w[:, :3],
self.pose_command_w[:, 3:],
self.robot.data.body_link_state_w[:, self.body_idx, :3],
self.robot.data.body_link_state_w[:, self.body_idx, 3:7],
self.robot.data.body_state_w[:, self.body_idx, :3],
self.robot.data.body_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 @@ class UniformPoseCommand(CommandTerm):
# -- goal pose
self.goal_pose_visualizer.visualize(self.pose_command_w[:, :3], self.pose_command_w[:, 3:])
# -- current body pose
body_link_state_w = self.robot.data.body_link_state_w[:, self.body_idx]
body_link_state_w = self.robot.data.body_state_w[:, self.body_idx]
self.current_pose_visualizer.visualize(body_link_state_w[:, :3], body_link_state_w[:, 3:7])
......@@ -114,10 +114,10 @@ class UniformVelocityCommand(CommandTerm):
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_com_lin_vel_b[:, :2], dim=-1) / max_command_step
torch.norm(self.vel_command_b[:, :2] - self.robot.data.root_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_com_ang_vel_b[:, 2]) / max_command_step
torch.abs(self.vel_command_b[:, 2] - self.robot.data.root_ang_vel_b[:, 2]) / max_command_step
)
def _resample_command(self, env_ids: Sequence[int]):
......@@ -184,11 +184,11 @@ class UniformVelocityCommand(CommandTerm):
return
# get marker location
# -- base state
base_pos_w = self.robot.data.root_link_pos_w.clone()
base_pos_w = self.robot.data.root_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_com_lin_vel_b[:, :2])
vel_arrow_scale, vel_arrow_quat = self._resolve_xy_velocity_to_arrow(self.robot.data.root_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 @@ class UniformVelocityCommand(CommandTerm):
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_link_quat_w
base_quat_w = self.robot.data.root_quat_w
arrow_quat = math_utils.quat_mul(base_quat_w, arrow_quat)
return arrow_scale, arrow_quat
......
......@@ -624,13 +624,13 @@ def push_by_setting_velocity(
asset: RigidObject | Articulation = env.scene[asset_cfg.name]
# velocities
vel_w = asset.data.root_com_vel_w[env_ids]
vel_w = asset.data.root_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_com_velocity_to_sim(vel_w, env_ids=env_ids)
asset.write_root_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_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)
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)
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_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)
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)
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_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)
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)
def reset_joints_by_scale(
......@@ -914,16 +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_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)
rigid_object.write_root_pose_to_sim(default_root_state[:, :7], env_ids=env_ids)
rigid_object.write_root_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_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)
articulation_asset.write_root_pose_to_sim(default_root_state[:, :7], env_ids=env_ids)
articulation_asset.write_root_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()
......
......@@ -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_link_pos_w[:, 2].unsqueeze(-1)
return asset.data.root_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_com_lin_vel_b
return asset.data.root_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_com_ang_vel_b
return asset.data.root_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_link_pos_w - env.scene.env_origins
return asset.data.root_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_link_quat_w
quat = asset.data.root_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_com_lin_vel_w
return asset.data.root_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_com_ang_vel_w
return asset.data.root_ang_vel_w
"""
......
......@@ -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_com_lin_vel_b[:, 2])
return torch.square(asset.data.root_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_com_ang_vel_b[:, :2]), dim=1)
return torch.sum(torch.square(asset.data.root_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_link_pos_w[:, 2] - adjusted_target_height)
return torch.square(asset.data.root_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_com_lin_vel_b[:, :2]),
torch.square(env.command_manager.get_command(command_name)[:, :2] - asset.data.root_lin_vel_b[:, :2]),
dim=1,
)
return torch.exp(-lin_vel_error / std**2)
......@@ -305,7 +305,5 @@ 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_com_ang_vel_b[:, 2]
)
ang_vel_error = torch.square(env.command_manager.get_command(command_name)[:, 2] - asset.data.root_ang_vel_b[:, 2])
return torch.exp(-ang_vel_error / std**2)
......@@ -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_link_pos_w[:, 2] < minimum_height
return asset.data.root_pos_w[:, 2] < minimum_height
"""
......
......@@ -161,7 +161,7 @@ class ViewportCameraController:
# 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_link_pos_w[self.cfg.env_index]
self.viewer_origin = self._env.scene[self.cfg.asset_name].data.root_pos_w[self.cfg.env_index]
# update the camera view
self.update_view_location()
......@@ -194,9 +194,7 @@ class ViewportCameraController:
# set origin type to asset_body
self.cfg.origin_type = "asset_body"
# update the camera origins
self.viewer_origin = (
self._env.scene[self.cfg.asset_name].data.body_link_pos_w[self.cfg.env_index, body_id].view(3)
)
self.viewer_origin = self._env.scene[self.cfg.asset_name].data.body_pos_w[self.cfg.env_index, body_id].view(3)
# update the camera view
self.update_view_location()
......
......@@ -364,10 +364,10 @@ class InteractiveScene:
state["articulation"] = dict()
for asset_name, articulation in self._articulations.items():
asset_state = dict()
asset_state["root_pose"] = articulation.data.root_link_state_w[:, :7].clone()
asset_state["root_pose"] = articulation.data.root_state_w[:, :7].clone()
if is_relative:
asset_state["root_pose"][:, :3] -= self.env_origins
asset_state["root_velocity"] = articulation.data.root_com_vel_w.clone()
asset_state["root_velocity"] = articulation.data.root_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 @@ class InteractiveScene:
state["rigid_object"] = dict()
for asset_name, rigid_object in self._rigid_objects.items():
asset_state = dict()
asset_state["root_pose"] = rigid_object.data.root_link_state_w[:, :7].clone()
asset_state["root_pose"] = rigid_object.data.root_state_w[:, :7].clone()
if is_relative:
asset_state["root_pose"][:, :3] -= self.env_origins
asset_state["root_velocity"] = rigid_object.data.root_com_vel_w.clone()
asset_state["root_velocity"] = rigid_object.data.root_vel_w.clone()
state["rigid_object"][asset_name] = asset_state
return state
......@@ -439,8 +439,8 @@ class InteractiveScene:
if is_relative:
root_pose[:, :3] += self.env_origins[env_ids]
root_velocity = asset_state["root_velocity"].clone()
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)
articulation.write_root_pose_to_sim(root_pose, env_ids=env_ids)
articulation.write_root_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 @@ class InteractiveScene:
if is_relative:
root_pose[:, :3] += self.env_origins[env_ids]
root_velocity = asset_state["root_velocity"].clone()
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)
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)
self.write_data_to_sim()
def write_data_to_sim(self):
......
......@@ -99,8 +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_link_pose_to_sim(root_state[:, :7])
robot.write_root_com_velocity_to_sim(root_state[:, 7:])
robot.write_root_pose_to_sim(root_state[:, :7])
robot.write_root_velocity_to_sim(root_state[:, 7:])
# reset dof state
joint_pos, joint_vel = robot.data.default_joint_pos, robot.data.default_joint_vel
robot.write_joint_state_to_sim(joint_pos, joint_vel)
......
......@@ -113,8 +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_link_pose_to_sim(root_state[:, :7])
robot.write_root_com_velocity_to_sim(root_state[:, 7:])
robot.write_root_pose_to_sim(root_state[:, :7])
robot.write_root_velocity_to_sim(root_state[:, 7:])
# joint state
joint_pos, joint_vel = robot.data.default_joint_pos.clone(), robot.data.default_joint_vel.clone()
robot.write_joint_state_to_sim(joint_pos, joint_vel)
......
......@@ -513,6 +513,8 @@ class TestRigidObjectCollection(unittest.TestCase):
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)])
env_ids = torch.tensor([x for x in range(num_envs)])
object_ids = torch.tensor([x for x in range(num_cubes)])
# Play sim
sim.reset()
......@@ -547,6 +549,8 @@ class TestRigidObjectCollection(unittest.TestCase):
# make quaternion a unit vector
rand_state[..., 3:7] = torch.nn.functional.normalize(rand_state[..., 3:7], dim=-1)
env_ids = env_ids.to(device)
object_ids = object_ids.to(device)
for i in range(10):
# perform step
......@@ -555,9 +559,19 @@ class TestRigidObjectCollection(unittest.TestCase):
cube_object.update(sim.cfg.dt)
if state_location == "com":
cube_object.write_object_com_state_to_sim(rand_state)
if i % 2 == 0:
cube_object.write_object_com_state_to_sim(rand_state)
else:
cube_object.write_object_com_state_to_sim(
rand_state, env_ids=env_ids, object_ids=object_ids
)
elif state_location == "link":
cube_object.write_object_link_state_to_sim(rand_state)
if i % 2 == 0:
cube_object.write_object_link_state_to_sim(rand_state)
else:
cube_object.write_object_link_state_to_sim(
rand_state, env_ids=env_ids, object_ids=object_ids
)
if state_location == "com":
torch.testing.assert_close(rand_state, cube_object.data.object_com_state_w)
......
......@@ -147,8 +147,8 @@ class TestDifferentialIKController(unittest.TestCase):
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_link_state_w[:, ee_frame_idx, 0:7]
root_pose_w = robot.data.root_link_state_w[:, 0:7]
ee_pose_w = robot.data.body_state_w[:, ee_frame_idx, 0:7]
root_pose_w = robot.data.root_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,10 +177,10 @@ class TestDifferentialIKController(unittest.TestCase):
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_link_state_w.clone()
root_state = robot.data.root_state_w.clone()
root_state[:, 3:7] = random_yaw_orientation(self.num_envs, self.sim.device)
robot.write_root_link_pose_to_sim(root_state[:, :7])
robot.write_root_com_velocity_to_sim(root_state[:, 7:])
robot.write_root_pose_to_sim(root_state[:, :7])
robot.write_root_velocity_to_sim(root_state[:, 7:])
robot.reset()
# reset actions
ee_pose_b_des[:] = self.ee_pose_b_des_set[current_goal_idx]
......@@ -195,8 +195,8 @@ class TestDifferentialIKController(unittest.TestCase):
# 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_link_state_w[:, ee_frame_idx, 0:7]
root_pose_w = robot.data.root_link_state_w[:, 0:7]
ee_pose_w = robot.data.body_state_w[:, ee_frame_idx, 0:7]
root_pose_w = robot.data.root_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, :])
......
......@@ -843,26 +843,24 @@ class TestOperationalSpaceController(unittest.TestCase):
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_link_quat_w))
root_rot_matrix = matrix_from_quat(quat_inv(robot.data.root_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_link_state_w[:, 0:7]
ee_pose_w = robot.data.body_link_state_w[:, ee_frame_idx, 0:7]
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]
)
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_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
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
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_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_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_vel_b = torch.cat([ee_lin_vel_b, ee_ang_vel_b], dim=-1)
# Calculate the contact force
......
......@@ -128,11 +128,11 @@ class CubeActionTerm(ActionTerm):
def apply_actions(self):
# implement a PD controller to track the target position
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
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
# set velocity targets
self._vel_command[:, :3] = self.p_gain * pos_error + self.d_gain * vel_error
self._asset.write_root_com_velocity_to_sim(self._vel_command)
self._asset.write_root_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_link_pos_w - env.scene.env_origins
return asset.data.root_pos_w - env.scene.env_origins
##
......
......@@ -98,8 +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_link_pose_to_sim(root_state[:, :7])
scene["robot"].write_root_com_velocity_to_sim(root_state[:, 7:])
scene["robot"].write_root_pose_to_sim(root_state[:, :7])
scene["robot"].write_root_velocity_to_sim(root_state[:, 7:])
# set joint positions with some noise
joint_pos, joint_vel = (
scene["robot"].data.default_joint_pos.clone(),
......
......@@ -130,13 +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_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_root_pose_to_sim(root_state[:, :7])
scene.articulations["robot_1"].write_root_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_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_root_pose_to_sim(root_state[:, :7])
scene.articulations["robot_2"].write_root_velocity_to_sim(root_state[:, 7:])
scene.articulations["robot_2"].write_joint_state_to_sim(joint_pos, joint_vel)
# reset buffers
scene.reset()
......
......@@ -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_link_pos_w.clone()
ball_initial_orientations = balls.data.root_link_quat_w.clone()
ball_initial_positions = balls.data.root_pos_w.clone()
ball_initial_orientations = balls.data.root_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_link_pose_to_sim(torch.cat([ball_initial_positions, ball_initial_orientations], dim=-1))
balls.write_root_pose_to_sim(torch.cat([ball_initial_positions, ball_initial_orientations], dim=-1))
balls.reset()
# reset the sensor
imu.reset()
......
......@@ -405,7 +405,7 @@ class TestContactSensor(unittest.TestCase):
duration = self.durations[idx]
while current_test_time < duration:
# set object states to contact the ground plane
shape.write_root_link_pose_to_sim(root_pose=test_pose)
shape.write_root_pose_to_sim(root_pose=test_pose)
# perform simulation step
self._perform_sim_step()
# increment contact time
......@@ -432,7 +432,7 @@ class TestContactSensor(unittest.TestCase):
dt=duration + self.sim_dt,
)
# switch the contact mode for 1 dt step before the next contact test begins.
shape.write_root_link_pose_to_sim(root_pose=reset_pose)
shape.write_root_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
......
......@@ -175,8 +175,8 @@ class TestFrameTransformer(unittest.TestCase):
joint_vel = scene.articulations["robot"].data.default_joint_vel
# -- set root state
# -- robot
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_root_pose_to_sim(root_state[:, :7])
scene.articulations["robot"].write_root_velocity_to_sim(root_state[:, 7:])
scene.articulations["robot"].write_joint_state_to_sim(joint_pos, joint_vel)
# reset buffers
scene.reset()
......@@ -193,9 +193,9 @@ class TestFrameTransformer(unittest.TestCase):
# check absolute frame transforms in world frame
# -- ground-truth
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]
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]
# -- 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
......@@ -276,8 +276,8 @@ class TestFrameTransformer(unittest.TestCase):
joint_vel = scene.articulations["robot"].data.default_joint_vel
# -- set root state
# -- robot
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_root_pose_to_sim(root_state[:, :7])
scene.articulations["robot"].write_root_velocity_to_sim(root_state[:, 7:])
scene.articulations["robot"].write_joint_state_to_sim(joint_pos, joint_vel)
# reset buffers
scene.reset()
......@@ -294,9 +294,9 @@ class TestFrameTransformer(unittest.TestCase):
# check absolute frame transforms in world frame
# -- ground-truth
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]
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]
# -- 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
......@@ -358,8 +358,8 @@ class TestFrameTransformer(unittest.TestCase):
joint_vel = scene.articulations["robot"].data.default_joint_vel
# -- set root state
# -- robot
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_root_pose_to_sim(root_state[:, :7])
scene.articulations["robot"].write_root_velocity_to_sim(root_state[:, 7:])
scene.articulations["robot"].write_joint_state_to_sim(joint_pos, joint_vel)
# reset buffers
scene.reset()
......@@ -376,9 +376,9 @@ class TestFrameTransformer(unittest.TestCase):
# check absolute frame transforms in world frame
# -- ground-truth
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]
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]
# -- 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
......@@ -450,8 +450,8 @@ class TestFrameTransformer(unittest.TestCase):
root_state[:, :3] += scene.env_origins
# -- set root state
# -- cube
scene["cube"].write_root_link_pose_to_sim(root_state[:, :7])
scene["cube"].write_root_com_velocity_to_sim(root_state[:, 7:])
scene["cube"].write_root_pose_to_sim(root_state[:, :7])
scene["cube"].write_root_velocity_to_sim(root_state[:, 7:])
# reset buffers
scene.reset()
......@@ -464,8 +464,8 @@ class TestFrameTransformer(unittest.TestCase):
# check absolute frame transforms in world frame
# -- ground-truth
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]
cube_pos_w_gt = scene["cube"].data.root_state_w[:, :3]
cube_quat_w_gt = scene["cube"].data.root_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
......@@ -538,8 +538,8 @@ class TestFrameTransformer(unittest.TestCase):
joint_vel = scene.articulations["robot"].data.default_joint_vel
# -- set root state
# -- robot
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_root_pose_to_sim(root_state[:, :7])
scene.articulations["robot"].write_root_velocity_to_sim(root_state[:, 7:])
scene.articulations["robot"].write_joint_state_to_sim(joint_pos, joint_vel)
# reset buffers
scene.reset()
......@@ -556,9 +556,9 @@ class TestFrameTransformer(unittest.TestCase):
# check absolute frame transforms in world frame
# -- ground-truth
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
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
# -- frame transformer
source_pos_w_tf = scene.sensors["frame_transformer"].data.source_pos_w
......
......@@ -189,12 +189,12 @@ class TestImu(unittest.TestCase):
for idx in range(200):
# set velocity
self.scene.rigid_objects["balls"].write_root_com_velocity_to_sim(
self.scene.rigid_objects["balls"].write_root_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_com_velocity_to_sim(
self.scene.rigid_objects["cube"].write_root_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 @@ class TestImu(unittest.TestCase):
)
# check the imu velocities
# NOTE: the expected lin_vel_b is the same as the set velocity, as write_root_com_velocity_to_sim is
# NOTE: the expected lin_vel_b is the same as the set velocity, as write_root_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 @@ class TestImu(unittest.TestCase):
"""Test the Imu sensor with a constant acceleration."""
for idx in range(100):
# set acceleration
self.scene.rigid_objects["balls"].write_root_com_velocity_to_sim(
self.scene.rigid_objects["balls"].write_root_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 @@ class TestImu(unittest.TestCase):
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_link_quat_w,
self.scene.rigid_objects["balls"].data.root_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 @@ class TestImu(unittest.TestCase):
# check the angular velocity
torch.testing.assert_close(
self.scene.sensors["imu_ball"].data.ang_vel_b,
self.scene.rigid_objects["balls"].data.root_com_ang_vel_b,
self.scene.rigid_objects["balls"].data.root_ang_vel_b,
rtol=1e-4,
atol=1e-4,
)
......@@ -438,7 +438,7 @@ class TestImu(unittest.TestCase):
# 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_com_velocity_to_sim(
self.scene.articulations["robot"].write_root_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 @@ class TestImu(unittest.TestCase):
for idx in range(10):
# set acceleration
self.scene.articulations["robot"].write_root_com_velocity_to_sim(
self.scene.articulations["robot"].write_root_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
)
......
......@@ -89,8 +89,8 @@ class AnymalCEnv(DirectRLEnv):
[
tensor
for tensor in (
self._robot.data.root_com_lin_vel_b,
self._robot.data.root_com_ang_vel_b,
self._robot.data.root_lin_vel_b,
self._robot.data.root_ang_vel_b,
self._robot.data.projected_gravity_b,
self._commands,
self._robot.data.joint_pos - self._robot.data.default_joint_pos,
......@@ -107,17 +107,15 @@ class AnymalCEnv(DirectRLEnv):
def _get_rewards(self) -> torch.Tensor:
# linear velocity tracking
lin_vel_error = torch.sum(
torch.square(self._commands[:, :2] - self._robot.data.root_com_lin_vel_b[:, :2]), dim=1
)
lin_vel_error = torch.sum(torch.square(self._commands[:, :2] - self._robot.data.root_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_com_ang_vel_b[:, 2])
yaw_rate_error = torch.square(self._commands[:, 2] - self._robot.data.root_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_com_lin_vel_b[:, 2])
z_vel_error = torch.square(self._robot.data.root_lin_vel_b[:, 2])
# angular velocity x/y
ang_vel_error = torch.sum(torch.square(self._robot.data.root_com_ang_vel_b[:, :2]), dim=1)
ang_vel_error = torch.sum(torch.square(self._robot.data.root_ang_vel_b[:, :2]), dim=1)
# joint torques
joint_torques = torch.sum(torch.square(self._robot.data.applied_torque), dim=1)
# joint acceleration
......@@ -180,8 +178,8 @@ class AnymalCEnv(DirectRLEnv):
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_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_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_joint_state_to_sim(joint_pos, joint_vel, None, env_ids)
# Logging
extras = dict()
......
......@@ -182,8 +182,8 @@ class CartDoublePendulumEnv(DirectMARLEnv):
self.joint_pos[env_ids] = joint_pos
self.joint_vel[env_ids] = joint_vel
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_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_joint_state_to_sim(joint_pos, joint_vel, None, env_ids)
......
......@@ -199,8 +199,8 @@ class CartpoleCameraEnv(DirectRLEnv):
self.joint_pos[env_ids] = joint_pos
self.joint_vel[env_ids] = joint_vel
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_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_joint_state_to_sim(joint_pos, joint_vel, None, env_ids)
......
......@@ -143,8 +143,8 @@ class CartpoleEnv(DirectRLEnv):
self.joint_pos[env_ids] = joint_pos
self.joint_vel[env_ids] = joint_vel
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_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_joint_state_to_sim(joint_pos, joint_vel, None, env_ids)
......
......@@ -189,18 +189,16 @@ class FactoryEnv(DirectRLEnv):
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_link_pos_w - self.scene.env_origins
self.fixed_quat = self._fixed_asset.data.root_link_quat_w
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.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.held_pos = self._held_asset.data.root_pos_w - self.scene.env_origins
self.held_quat = self._held_asset.data.root_quat_w
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]
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]
jacobians = self._robot.root_physx_view.get_jacobians()
......@@ -554,15 +552,15 @@ class FactoryEnv(DirectRLEnv):
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_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.write_root_pose_to_sim(held_state[:, 0:7], env_ids=env_ids)
self._held_asset.write_root_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_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.write_root_pose_to_sim(fixed_state[:, 0:7], env_ids=env_ids)
self._fixed_asset.write_root_velocity_to_sim(fixed_state[:, 7:], env_ids=env_ids)
self._fixed_asset.reset()
def set_pos_inverse_kinematics(self, env_ids):
......@@ -685,8 +683,8 @@ class FactoryEnv(DirectRLEnv):
# (1.c.) Velocity
fixed_state[:, 7:] = 0.0 # vel
# (1.d.) Update values.
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.write_root_pose_to_sim(fixed_state[:, 0:7], env_ids=env_ids)
self._fixed_asset.write_root_velocity_to_sim(fixed_state[:, 7:], env_ids=env_ids)
self._fixed_asset.reset()
# (1.e.) Noisy position observation.
......@@ -772,15 +770,15 @@ class FactoryEnv(DirectRLEnv):
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_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.write_root_pose_to_sim(small_gear_state[:, 0:7], env_ids=env_ids)
self._small_gear_asset.write_root_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_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.write_root_pose_to_sim(large_gear_state[:, 0:7], env_ids=env_ids)
self._large_gear_asset.write_root_velocity_to_sim(large_gear_state[:, 7:], env_ids=env_ids)
self._large_gear_asset.reset()
# (3) Randomize asset-in-gripper location.
......@@ -822,8 +820,8 @@ class FactoryEnv(DirectRLEnv):
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_link_pose_to_sim(held_state[:, 0:7])
self._held_asset.write_root_com_velocity_to_sim(held_state[:, 7:])
self._held_asset.write_root_pose_to_sim(held_state[:, 0:7])
self._held_asset.write_root_velocity_to_sim(held_state[:, 7:])
self._held_asset.reset()
# Close hand
......
......@@ -299,8 +299,8 @@ class FrankaCabinetEnv(DirectRLEnv):
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_link_pos_w[:, self.left_finger_link_idx]
robot_right_finger_pos = self._robot.data.body_link_pos_w[:, self.right_finger_link_idx]
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]
return self._compute_rewards(
self.actions,
......@@ -372,10 +372,10 @@ class FrankaCabinetEnv(DirectRLEnv):
if env_ids is None:
env_ids = self._robot._ALL_INDICES
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]
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]
(
self.robot_grasp_rot[env_ids],
self.robot_grasp_pos[env_ids],
......
......@@ -221,8 +221,8 @@ class InHandManipulationEnv(DirectRLEnv):
)
object_default_state[:, 7:] = torch.zeros_like(self.object.data.default_root_state[env_ids, 7:])
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)
self.object.write_root_pose_to_sim(object_default_state[:, :7], env_ids)
self.object.write_root_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]
......@@ -261,22 +261,22 @@ class InHandManipulationEnv(DirectRLEnv):
def _compute_intermediate_values(self):
# data for hand
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.hand.data.body_pos_w[:, self.finger_bodies]
self.fingertip_rot = self.hand.data.body_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_com_vel_w[:, self.finger_bodies]
self.fingertip_velocities = self.hand.data.body_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_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
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
def compute_reduced_observations(self):
# Per https://arxiv.org/pdf/1808.00177.pdf Table 2
......
......@@ -68,8 +68,8 @@ class LocomotionEnv(DirectRLEnv):
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_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.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.dof_pos, self.dof_vel = self.robot.data.joint_pos, self.robot.data.joint_vel
(
......@@ -161,8 +161,8 @@ class LocomotionEnv(DirectRLEnv):
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_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_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_joint_state_to_sim(joint_pos, joint_vel, None, env_ids)
to_target = self.targets[env_ids] - default_root_state[:, :3]
......
......@@ -154,12 +154,12 @@ class QuadcopterEnv(DirectRLEnv):
def _get_observations(self) -> dict:
desired_pos_b, _ = subtract_frame_transforms(
self._robot.data.root_link_state_w[:, :3], self._robot.data.root_link_state_w[:, 3:7], self._desired_pos_w
self._robot.data.root_state_w[:, :3], self._robot.data.root_state_w[:, 3:7], self._desired_pos_w
)
obs = torch.cat(
[
self._robot.data.root_com_lin_vel_b,
self._robot.data.root_com_ang_vel_b,
self._robot.data.root_lin_vel_b,
self._robot.data.root_ang_vel_b,
self._robot.data.projected_gravity_b,
desired_pos_b,
],
......@@ -169,9 +169,9 @@ class QuadcopterEnv(DirectRLEnv):
return observations
def _get_rewards(self) -> torch.Tensor:
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)
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)
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,9 +186,7 @@ class QuadcopterEnv(DirectRLEnv):
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_link_pos_w[:, 2] < 0.1, self._robot.data.root_link_pos_w[:, 2] > 2.0
)
died = torch.logical_or(self._robot.data.root_pos_w[:, 2] < 0.1, self._robot.data.root_pos_w[:, 2] > 2.0)
return died, time_out
def _reset_idx(self, env_ids: torch.Tensor | None):
......@@ -197,7 +195,7 @@ class QuadcopterEnv(DirectRLEnv):
# Logging
final_distance_to_goal = torch.linalg.norm(
self._desired_pos_w[env_ids] - self._robot.data.root_link_pos_w[env_ids], dim=1
self._desired_pos_w[env_ids] - self._robot.data.root_pos_w[env_ids], dim=1
).mean()
extras = dict()
for key in self._episode_sums.keys():
......@@ -228,8 +226,8 @@ class QuadcopterEnv(DirectRLEnv):
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_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_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_joint_state_to_sim(joint_pos, joint_vel, None, env_ids)
def _set_debug_vis_impl(self, debug_vis: bool):
......
......@@ -322,8 +322,8 @@ class ShadowHandOverEnv(DirectMARLEnv):
)
object_default_state[:, 7:] = torch.zeros_like(self.object.data.default_root_state[env_ids, 7:])
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)
self.object.write_root_pose_to_sim(object_default_state[:, :7], env_ids)
self.object.write_root_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]
......@@ -377,33 +377,33 @@ class ShadowHandOverEnv(DirectMARLEnv):
def _compute_intermediate_values(self):
# data for right hand
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.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.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_com_vel_w[:, self.finger_bodies]
self.right_fingertip_velocities = self.right_hand.data.body_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_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.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.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_com_vel_w[:, self.finger_bodies]
self.left_fingertip_velocities = self.left_hand.data.body_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_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
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
@torch.jit.script
......
......@@ -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_link_quat_w)
roll, _, yaw = math_utils.euler_xyz_from_quat(asset.data.root_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_link_pos_w[:, :3]
to_target_pos = torch.tensor(target_pos, device=env.device) - asset.data.root_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_link_quat_w, asset.data.FORWARD_VEC_B)
heading_vec = math_utils.quat_rotate(asset.data.root_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_link_pos_w[:, :3]
to_target_pos = torch.tensor(target_pos, device=env.device) - asset.data.root_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_link_quat_w)
_, _, yaw = math_utils.euler_xyz_from_quat(asset.data.root_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))
......
......@@ -53,7 +53,7 @@ class progress_reward(ManagerTermBase):
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_link_pos_w[env_ids, :3]
to_target_pos = target_pos - asset.data.root_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 @@ class progress_reward(ManagerTermBase):
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_link_pos_w[:, :3]
to_target_pos = target_pos - asset.data.root_pos_w[:, :3]
to_target_pos[:, 2] = 0.0
# update history buffer and compute new potential
self.prev_potentials[:] = self.potentials[:]
......
......@@ -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_com_lin_vel_b[:, :2], 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)
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_com_ang_vel_b[:, 2]).unsqueeze(1), dim=1)
ang_vel_error = torch.linalg.norm((target - asset.data.root_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_com_lin_vel_b[:, :2]), dim=1)
lin_vel_error = torch.linalg.norm((target - asset.data.root_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 @@ class GaitReward(ManagerTermBase):
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_com_lin_vel_b[:, :2], dim=1)
body_vel = torch.linalg.norm(self.asset.data.root_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,10 +182,8 @@ 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_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)
)
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))
reward = foot_z_target_error * foot_velocity_tanh
return torch.exp(-torch.sum(reward, dim=1) / std)
......@@ -219,8 +217,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_com_lin_vel_b[:, 2]) + 0.2 * torch.sum(
torch.abs(asset.data.root_com_ang_vel_b[:, :2]), dim=1
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
)
......@@ -245,7 +243,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_com_lin_vel_w[:, asset_cfg.body_ids, :2], dim=2)
foot_planar_velocity = torch.linalg.norm(asset.data.body_lin_vel_w[:, asset_cfg.body_ids, :2], dim=2)
reward = is_contact * foot_planar_velocity
return torch.sum(reward, dim=1)
......@@ -265,7 +263,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_com_lin_vel_b[:, :2], dim=1)
body_vel = torch.linalg.norm(asset.data.root_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)
......
......@@ -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_link_pos_w[env_ids, :2] - env.scene.env_origins[env_ids, :2], dim=1)
distance = torch.norm(asset.data.root_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
......
......@@ -78,7 +78,7 @@ def feet_slide(env, sensor_cfg: SceneEntityCfg, asset_cfg: SceneEntityCfg = Scen
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_com_lin_vel_w[:, asset_cfg.body_ids, :2]
body_vel = asset.data.body_lin_vel_w[:, asset_cfg.body_ids, :2]
reward = torch.sum(body_vel.norm(dim=-1) * contacts, dim=1)
return reward
......@@ -89,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_link_quat_w), asset.data.root_com_lin_vel_w[:, :3])
vel_yaw = quat_rotate_inverse(yaw_quat(asset.data.root_quat_w), asset.data.root_lin_vel_w[:, :3])
lin_vel_error = torch.sum(
torch.square(env.command_manager.get_command(command_name)[:, :2] - vel_yaw[:, :2]), dim=1
)
......@@ -102,7 +102,5 @@ 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_com_ang_vel_w[:, 2]
)
ang_vel_error = torch.square(env.command_manager.get_command(command_name)[:, 2] - asset.data.root_ang_vel_w[:, 2])
return torch.exp(-ang_vel_error / std**2)
......@@ -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_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
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
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'.")
......@@ -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_link_pos_w - ee_tf_data.target_pos_w[..., 0, :]
return object_data.root_pos_w - ee_tf_data.target_pos_w[..., 0, :]
def rel_ee_drawer_distance(env: ManagerBasedRLEnv) -> torch.Tensor:
......
......@@ -95,10 +95,10 @@ class InHandReOrientationCommand(CommandTerm):
# logs data
# -- compute the orientation error
self.metrics["orientation_error"] = math_utils.quat_error_magnitude(
self.object.data.root_link_quat_w, self.quat_command_w
self.object.data.root_quat_w, self.quat_command_w
)
# -- compute the position error
self.metrics["position_error"] = torch.norm(self.object.data.root_link_pos_w - self.pos_command_w, dim=1)
self.metrics["position_error"] = torch.norm(self.object.data.root_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()
......
......@@ -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_link_quat_w
asset_quat_w = asset.data.root_quat_w
# compute quaternion difference
quat = math_utils.quat_mul(asset_quat_w, math_utils.quat_conjugate(goal_quat_w))
......
......@@ -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_link_quat_w, goal_quat_w)
dtheta = math_utils.quat_error_magnitude(asset.data.root_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_link_pos_w - env.scene.env_origins
object_pos_e = asset.data.root_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_link_quat_w, goal_quat_w)
dtheta = math_utils.quat_error_magnitude(asset.data.root_quat_w, goal_quat_w)
return 1.0 / (dtheta + rot_eps)
......@@ -50,7 +50,7 @@ def object_away_from_goal(
asset = env.scene[object_cfg.name]
# object pos
asset_pos_e = asset.data.root_link_pos_w - env.scene.env_origins
asset_pos_e = asset.data.root_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_link_pos_w - object.data.root_link_pos_w, dim=1)
dist = torch.norm(robot.data.root_pos_w - object.data.root_pos_w, dim=1)
return dist > threshold
......@@ -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_link_pos_w[:, :3]
object_pos_w = object.data.root_pos_w[:, :3]
object_pos_b, _ = subtract_frame_transforms(
robot.data.root_link_state_w[:, :3], robot.data.root_link_state_w[:, 3:7], object_pos_w
robot.data.root_state_w[:, :3], robot.data.root_state_w[:, 3:7], object_pos_w
)
return object_pos_b
......@@ -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_link_pos_w[:, 2] > minimal_height, 1.0, 0.0)
return torch.where(object.data.root_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_link_pos_w
cube_pos_w = object.data.root_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,10 +60,8 @@ def object_goal_distance(
command = env.command_manager.get_command(command_name)
# compute the desired position in the world frame
des_pos_b = command[:, :3]
des_pos_w, _ = combine_frame_transforms(
robot.data.root_link_state_w[:, :3], robot.data.root_link_state_w[:, 3:7], des_pos_b
)
des_pos_w, _ = combine_frame_transforms(robot.data.root_state_w[:, :3], robot.data.root_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_link_pos_w[:, :3], dim=1)
distance = torch.norm(des_pos_w - object.data.root_pos_w[:, :3], dim=1)
# rewarded if the object is lifted above the threshold
return (object.data.root_link_pos_w[:, 2] > minimal_height) * (1 - torch.tanh(distance / std))
return (object.data.root_pos_w[:, 2] > minimal_height) * (1 - torch.tanh(distance / std))
......@@ -45,11 +45,9 @@ 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_link_state_w[:, :3], robot.data.root_link_state_w[:, 3:7], des_pos_b
)
des_pos_w, _ = combine_frame_transforms(robot.data.root_state_w[:, :3], robot.data.root_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_link_pos_w[:, :3], dim=1)
distance = torch.norm(des_pos_w - object.data.root_pos_w[:, :3], dim=1)
# rewarded if the object is lifted above the threshold
return distance < threshold
......@@ -28,10 +28,8 @@ def position_command_error(env: ManagerBasedRLEnv, command_name: str, asset_cfg:
command = env.command_manager.get_command(command_name)
# obtain the desired and current positions
des_pos_b = command[:, :3]
des_pos_w, _ = combine_frame_transforms(
asset.data.root_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
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
return torch.norm(curr_pos_w - des_pos_w, dim=1)
......@@ -48,10 +46,8 @@ def position_command_error_tanh(
command = env.command_manager.get_command(command_name)
# obtain the desired and current positions
des_pos_b = command[:, :3]
des_pos_w, _ = combine_frame_transforms(
asset.data.root_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
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
distance = torch.norm(curr_pos_w - des_pos_w, dim=1)
return 1 - torch.tanh(distance / std)
......@@ -68,6 +64,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_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
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
return quat_error_magnitude(curr_quat_w, des_quat_w)
......@@ -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_link_pose_to_sim(
asset.write_root_pose_to_sim(
torch.cat([positions, orientations], dim=-1), env_ids=torch.tensor([cur_env], device=env.device)
)
asset.write_root_com_velocity_to_sim(
asset.write_root_velocity_to_sim(
torch.zeros(1, 6, device=env.device), env_ids=torch.tensor([cur_env], device=env.device)
)
......
......@@ -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_link_pos_w, cube_2.data.root_link_pos_w, cube_3.data.root_link_pos_w), dim=1)
return torch.cat((cube_1.data.root_pos_w, cube_2.data.root_pos_w, cube_3.data.root_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_link_quat_w, cube_2.data.root_link_quat_w, cube_3.data.root_link_quat_w), dim=1)
return torch.cat((cube_1.data.root_quat_w, cube_2.data.root_quat_w, cube_3.data.root_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_link_pos_w
cube_1_quat_w = cube_1.data.root_link_quat_w
cube_1_pos_w = cube_1.data.root_pos_w
cube_1_quat_w = cube_1.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_2_pos_w = cube_2.data.root_pos_w
cube_2_quat_w = cube_2.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
cube_3_pos_w = cube_3.data.root_pos_w
cube_3_quat_w = cube_3.data.root_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_link_pos_w
object_pos = object.data.root_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_link_pos_w - lower_object.data.root_link_pos_w
pos_diff = upper_object.data.root_pos_w - lower_object.data.root_pos_w
height_dist = torch.linalg.vector_norm(pos_diff[:, 2:], dim=1)
xy_dist = torch.linalg.vector_norm(pos_diff[:, :2], dim=1)
......
......@@ -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_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
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
# Compute cube position difference in x-y plane
xy_dist_c12 = torch.norm(pos_diff_c12[:, :2], dim=1)
......
......@@ -134,11 +134,11 @@ class PreTrainedPolicyAction(ActionTerm):
return
# get marker location
# -- base state
base_pos_w = self.robot.data.root_link_pos_w.clone()
base_pos_w = self.robot.data.root_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_com_lin_vel_b[:, :2])
vel_arrow_scale, vel_arrow_quat = self._resolve_xy_velocity_to_arrow(self.robot.data.root_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)
......@@ -159,7 +159,7 @@ class PreTrainedPolicyAction(ActionTerm):
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_link_quat_w
base_quat_w = self.robot.data.root_quat_w
arrow_quat = math_utils.quat_mul(base_quat_w, arrow_quat)
return arrow_scale, arrow_quat
......
......@@ -115,8 +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_link_pose_to_sim(root_state[:, :7])
robot.write_root_com_velocity_to_sim(root_state[:, 7:])
robot.write_root_pose_to_sim(root_state[:, :7])
robot.write_root_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
......
......@@ -179,8 +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_link_pose_to_sim(root_state[:, :7])
robot.write_root_com_velocity_to_sim(root_state[:, 7:])
robot.write_root_pose_to_sim(root_state[:, :7])
robot.write_root_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)
......
......@@ -97,8 +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_link_pose_to_sim(root_state[:, :7])
robot.write_root_com_velocity_to_sim(root_state[:, 7:])
robot.write_root_pose_to_sim(root_state[:, :7])
robot.write_root_velocity_to_sim(root_state[:, 7:])
robot.reset()
# reset command
print(">>>>>>>> Reset!")
......
......@@ -187,8 +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_link_pose_to_sim(root_state[:, :7])
scene["robot"].write_root_com_velocity_to_sim(root_state[:, 7:])
scene["robot"].write_root_pose_to_sim(root_state[:, :7])
scene["robot"].write_root_velocity_to_sim(root_state[:, 7:])
# set joint positions with some noise
joint_pos, joint_vel = (
scene["robot"].data.default_joint_pos.clone(),
......
......@@ -113,8 +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_link_pose_to_sim(root_state[:, :7])
robot.write_root_com_velocity_to_sim(root_state[:, 7:])
robot.write_root_pose_to_sim(root_state[:, :7])
robot.write_root_velocity_to_sim(root_state[:, 7:])
# joint state
joint_pos, joint_vel = robot.data.default_joint_pos.clone(), robot.data.default_joint_vel.clone()
robot.write_joint_state_to_sim(joint_pos, joint_vel)
......
......@@ -240,8 +240,8 @@ 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_link_pose_to_sim(root_state[:, :7])
rigid_object.write_root_com_velocity_to_sim(root_state[:, 7:])
rigid_object.write_root_pose_to_sim(root_state[:, :7])
rigid_object.write_root_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)
......@@ -251,8 +251,8 @@ def run_simulator(sim: SimulationContext, scene: InteractiveScene):
# -- root state
root_state = robot.data.default_root_state.clone()
root_state[:, :3] += scene.env_origins
robot.write_root_link_pose_to_sim(root_state[:, :7])
robot.write_root_com_velocity_to_sim(root_state[:, 7:])
robot.write_root_pose_to_sim(root_state[:, :7])
robot.write_root_velocity_to_sim(root_state[:, 7:])
# -- joint state
joint_pos, joint_vel = robot.data.default_joint_pos.clone(), robot.data.default_joint_vel.clone()
robot.write_joint_state_to_sim(joint_pos, joint_vel)
......
......@@ -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_link_pose_to_sim(robot.data.default_root_state[:, :7])
robot.write_root_com_velocity_to_sim(robot.data.default_root_state[:, 7:])
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.reset()
# reset command
print(">>>>>>>> Reset!")
......
......@@ -142,8 +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_link_pose_to_sim(root_state[:, :7])
robot.write_root_com_velocity_to_sim(root_state[:, 7:])
robot.write_root_pose_to_sim(root_state[:, :7])
robot.write_root_velocity_to_sim(root_state[:, 7:])
# joint state
joint_pos, joint_vel = robot.data.default_joint_pos.clone(), robot.data.default_joint_vel.clone()
robot.write_joint_state_to_sim(joint_pos, joint_vel)
......
......@@ -109,8 +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_link_pose_to_sim(root_state[:, :7])
scene["robot"].write_root_com_velocity_to_sim(root_state[:, 7:])
scene["robot"].write_root_pose_to_sim(root_state[:, :7])
scene["robot"].write_root_velocity_to_sim(root_state[:, 7:])
# set joint positions with some noise
joint_pos, joint_vel = (
scene["robot"].data.default_joint_pos.clone(),
......
......@@ -105,8 +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_link_pose_to_sim(root_state[:, :7])
scene["robot"].write_root_com_velocity_to_sim(root_state[:, 7:])
scene["robot"].write_root_pose_to_sim(root_state[:, :7])
scene["robot"].write_root_velocity_to_sim(root_state[:, 7:])
# set joint positions with some noise
joint_pos, joint_vel = (
scene["robot"].data.default_joint_pos.clone(),
......
......@@ -93,8 +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_link_pose_to_sim(root_state[:, :7])
scene["robot"].write_root_com_velocity_to_sim(root_state[:, 7:])
scene["robot"].write_root_pose_to_sim(root_state[:, :7])
scene["robot"].write_root_velocity_to_sim(root_state[:, 7:])
# set joint positions with some noise
joint_pos, joint_vel = (
scene["robot"].data.default_joint_pos.clone(),
......
......@@ -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_link_pos_w - env.unwrapped.scene.env_origins
object_position = object_data.root_pos_w - env.unwrapped.scene.env_origins
# -- target object frame
desired_position = env.unwrapped.command_manager.get_command("object_pose")[..., :3]
......
......@@ -94,8 +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_link_pose_to_sim(root_state[:, :7])
robot.write_root_com_velocity_to_sim(root_state[:, 7:])
robot.write_root_pose_to_sim(root_state[:, :7])
robot.write_root_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
......
......@@ -103,8 +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_link_pose_to_sim(root_state[:, :7])
cone_object.write_root_com_velocity_to_sim(root_state[:, 7:])
cone_object.write_root_pose_to_sim(root_state[:, :7])
cone_object.write_root_velocity_to_sim(root_state[:, 7:])
# reset buffers
cone_object.reset()
print("----------------------------------------")
......@@ -120,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_link_state_w[:, :3]}")
print(f"Root position (in world): {cone_object.data.root_state_w[:, :3]}")
def main():
......
......@@ -83,8 +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_link_pose_to_sim(root_state[:, :7])
robot.write_root_com_velocity_to_sim(root_state[:, 7:])
robot.write_root_pose_to_sim(root_state[:, :7])
robot.write_root_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
......
......@@ -120,11 +120,11 @@ class CubeActionTerm(ActionTerm):
def apply_actions(self):
# implement a PD controller to track the target position
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
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
# set velocity targets
self._vel_command[:, :3] = self.p_gain * pos_error + self.d_gain * vel_error
self._asset.write_root_com_velocity_to_sim(self._vel_command)
self._asset.write_root_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_link_pos_w - env.scene.env_origins
return asset.data.root_pos_w - env.scene.env_origins
##
......
......@@ -113,8 +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_link_pose_to_sim(root_state[:, :7])
scene["robot"].write_root_com_velocity_to_sim(root_state[:, 7:])
scene["robot"].write_root_pose_to_sim(root_state[:, :7])
scene["robot"].write_root_velocity_to_sim(root_state[:, 7:])
# set joint positions with some noise
joint_pos, joint_vel = (
scene["robot"].data.default_joint_pos.clone(),
......
......@@ -109,8 +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_link_pose_to_sim(ball_default_state[:, :7])
balls.write_root_com_velocity_to_sim(ball_default_state[:, 7:])
balls.write_root_pose_to_sim(ball_default_state[:, :7])
balls.write_root_velocity_to_sim(ball_default_state[:, 7:])
# reset the sensor
ray_caster.reset()
# reset the counter
......
......@@ -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_link_state_w[:, robot_entity_cfg.body_ids[0], 0:7]
root_pose_w = robot.data.root_link_state_w[:, 0:7]
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]
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_link_state_w[:, robot_entity_cfg.body_ids[0], 0:7]
ee_pose_w = robot.data.body_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])
......
......@@ -318,26 +318,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_link_quat_w))
root_rot_matrix = matrix_from_quat(quat_inv(robot.data.root_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_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]
root_pos_w = robot.data.root_pos_w
root_quat_w = robot.data.root_quat_w
ee_pos_w = robot.data.body_pos_w[:, ee_frame_idx]
ee_quat_w = robot.data.body_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_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
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
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_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_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_vel_b = torch.cat([ee_lin_vel_b, ee_ang_vel_b], dim=-1)
# Calculate the contact force
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment