Unverified Commit 5d86a8bf authored by James Tigue's avatar James Tigue Committed by GitHub

Refactors pose and velocities to link frame and COM frame APIs (#966)

# Description

Currently the root_physx_views of the rigid bodies and articulations
output linear and angular velocities of the com of bodies rather than
the link frame. This PR transforms the velocities and accelerations to
the link frame of the body.

Fixes #942 

## Type of change

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

## 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
- [ ] 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 avatarjtigue-bdai <166445701+jtigue-bdai@users.noreply.github.com>
Signed-off-by: 's avatarKelly Guo <kellyg@nvidia.com>
Co-authored-by: 's avatarDavid Hoeller <dhoeller@nvidia.com>
Co-authored-by: 's avatarKelly Guo <kellyg@nvidia.com>
Co-authored-by: 's avatarKelly Guo <kellyguo123@hotmail.com>
parent 019e3523
......@@ -814,9 +814,9 @@ The ``progress_buf`` variable has also been renamed to ``episode_length_buf``.
| | |
| | self.joint_pos[env_ids] = joint_pos |
| | |
| | self.cartpole.write_root_pose_to_sim( |
| | self.cartpole.write_root_link_pose_to_sim( |
| | default_root_state[:, :7], env_ids) |
| | self.cartpole.write_root_velocity_to_sim( |
| | self.cartpole.write_root_com_velocity_to_sim( |
| | default_root_state[:, 7:], env_ids) |
| | self.cartpole.write_joint_state_to_sim( |
| | joint_pos, joint_vel, None, env_ids) |
......
......@@ -364,8 +364,8 @@ In Isaac Lab, ``root_pose`` and ``root_velocity`` have been combined into single
.. code-block::python
self.cartpole.write_root_pose_to_sim(default_root_state[:, :7], env_ids)
self.cartpole.write_root_velocity_to_sim(default_root_state[:, 7:], env_ids)
self.cartpole.write_root_link_pose_to_sim(default_root_state[:, :7], env_ids)
self.cartpole.write_root_com_velocity_to_sim(default_root_state[:, 7:], env_ids)
Creating a New Environment
......@@ -738,9 +738,9 @@ reset the ``episode_length_buf`` buffer.
| 1.0 - 2.0 * torch.rand(num_resets, device=self._device)) | self.joint_pos[env_ids] = joint_pos |
| | self.joint_vel[env_ids] = joint_vel |
| # apply resets | |
| indices = env_ids.to(dtype=torch.int32) | self.cartpole.write_root_pose_to_sim( |
| indices = env_ids.to(dtype=torch.int32) | self.cartpole.write_root_link_pose_to_sim( |
| self._cartpoles.set_joint_positions(dof_pos, indices=indices) | default_root_state[:, :7], env_ids) |
| self._cartpoles.set_joint_velocities(dof_vel, indices=indices) | self.cartpole.write_root_velocity_to_sim( |
| self._cartpoles.set_joint_velocities(dof_vel, indices=indices) | self.cartpole.write_root_com_velocity_to_sim( |
| | default_root_state[:, 7:], env_ids) |
| # bookkeeping | self.cartpole.write_joint_state_to_sim( |
| self.reset_buf[env_ids] = 0 | joint_pos, joint_vel, None, env_ids) |
......
......@@ -42,6 +42,7 @@ Using a ray caster sensor requires a **pattern** and a parent xform to be attach
update_period = 1/60,
offset=RayCasterCfg.OffsetCfg(pos=(0,0,0.5)),
mesh_prim_paths=["/World/Ground"],
attach_yaw_only=True,
pattern_cfg=patterns.LidarPatternCfg(
channels = 100,
vertical_fov_range=[-90, 90],
......
......@@ -66,8 +66,8 @@ Similar to a rigid object, an articulation also has a root state. This state cor
articulation tree. On top of the root state, an articulation also has joint states. These states correspond to the
joint positions and velocities.
To reset the articulation, we first set the root state by calling the :meth:`Articulation.write_root_state_to_sim`
method. Similarly, we set the joint states by calling the :meth:`Articulation.write_joint_state_to_sim` method.
To reset the articulation, we first set the root state by calling the :meth:`Articulation.write_root_link_pose_to_sim` and :meth:`Articulation.write_root_com_velocity_to_sim`
methods. Similarly, we set the joint states by calling the :meth:`Articulation.write_joint_state_to_sim` method.
Finally, we call the :meth:`Articulation.reset` method to reset any internal buffers and caches.
.. literalinclude:: ../../../../source/standalone/tutorials/01_assets/run_articulation.py
......
......@@ -149,7 +149,7 @@ the average position of all the nodes in the mesh.
.. literalinclude:: ../../../../source/standalone/tutorials/01_assets/run_deformable_object.py
:language: python
:start-at: # update buffers
:end-at: print(f"Root position (in world): {cube_object.data.root_pos_w[:, :3]}")
:end-at: print(f"Root position (in world): {cube_object.data.root_link_pos_w[:, :3]}")
The Code Execution
......
......@@ -96,7 +96,7 @@ desired state of the rigid object prim into the world frame before setting it.
We use the :attr:`assets.RigidObject.data.default_root_state` attribute to get the default root state of the
spawned rigid object prims. This default state can be configured from the :attr:`assets.RigidObjectCfg.init_state`
attribute, which we left as identity in this tutorial. We then randomize the translation of the root state and
set the desired state of the rigid object prim using the :meth:`assets.RigidObject.write_root_state_to_sim` method.
set the desired state of the rigid object prim using the :meth:`assets.RigidObject.write_root_link_pose_to_sim` and :meth:`assets.RigidObject.write_root_com_velocity_to_sim` methods.
As the name suggests, this method writes the root state of the rigid object prim into the simulation buffer.
.. literalinclude:: ../../../../source/standalone/tutorials/01_assets/run_rigid_object.py
......
......@@ -79,7 +79,7 @@ joint positions, current end-effector pose, and the Jacobian matrix.
While the attribute :attr:`assets.ArticulationData.joint_pos` provides the joint positions,
we only want the joint positions of the robot's arm, and not the gripper. Similarly, while
the attribute :attr:`assets.ArticulationData.body_state_w` provides the state of all the
the attribute :attr:`assets.ArticulationData.body_link_state_w` provides the state of all the
robot's bodies, we only want the state of the robot's end-effector. Thus, we need to
index into these arrays to obtain the desired quantities.
......
[package]
# Note: Semantic Versioning is used: https://semver.org/
version = "0.29.3"
version = "0.30.0"
# Description
title = "Isaac Lab framework for Robot Learning"
......
Changelog
---------
0.30.0 (2024-12-16)
~~~~~~~~~~~~~~~~~~~
Changed
^^^^^^^
* Previously, physx returns the rigid bodies and articulations velocities in the com of bodies rather than the link frame, while poses are in link frames. We now explicitly provide :attr:`body_link_state` and :attr:`body_com_state` APIs replacing the previous :attr:`body_state` API. Previous APIs are now marked as deprecated. Please update any code using the previous pose and velocity APIs to use the new ``*_link_*`` or ``*_com_*`` APIs in :attr:`omni.isaac_lab.assets.RigidBody`, :attr:`omni.isaac_lab.assets.RigidBodyCollection`, and :attr:`omni.isaac_lab.assets.Articulation`.
0.29.3 (2024-12-16)
~~~~~~~~~~~~~~~~~~~
......
......@@ -150,7 +150,7 @@ class NonHolonomicAction(ActionTerm):
def apply_actions(self):
# obtain current heading
quat_w = self._asset.data.body_quat_w[:, self._body_idx].view(self.num_envs, 4)
quat_w = self._asset.data.body_link_quat_w[:, self._body_idx].view(self.num_envs, 4)
yaw_w = euler_xyz_from_quat(quat_w)[2]
# compute joint velocities targets
self._joint_vel_command[:, 0] = torch.cos(yaw_w) * self.processed_actions[:, 0] # x
......
......@@ -142,7 +142,7 @@ class DifferentialInverseKinematicsAction(ActionTerm):
@property
def jacobian_b(self) -> torch.Tensor:
jacobian = self.jacobian_w
base_rot = self._asset.data.root_quat_w
base_rot = self._asset.data.root_link_quat_w
base_rot_matrix = math_utils.matrix_from_quat(math_utils.quat_inv(base_rot))
jacobian[:, :3, :] = torch.bmm(base_rot_matrix, jacobian[:, :3, :])
jacobian[:, 3:, :] = torch.bmm(base_rot_matrix, jacobian[:, 3:, :])
......@@ -192,12 +192,12 @@ class DifferentialInverseKinematicsAction(ActionTerm):
A tuple of the body's position and orientation in the root frame.
"""
# obtain quantities from simulation
ee_pose_w = self._asset.data.body_state_w[:, self._body_idx, :7]
root_pose_w = self._asset.data.root_state_w[:, :7]
ee_pos_w = self._asset.data.body_link_pos_w[:, self._body_idx]
ee_quat_w = self._asset.data.body_link_quat_w[:, self._body_idx]
root_pos_w = self._asset.data.root_link_pos_w
root_quat_w = self._asset.data.root_link_quat_w
# compute the pose of the body in the root frame
ee_pose_b, ee_quat_b = math_utils.subtract_frame_transforms(
root_pose_w[:, 0:3], root_pose_w[:, 3:7], ee_pose_w[:, 0:3], ee_pose_w[:, 3:7]
)
ee_pose_b, ee_quat_b = math_utils.subtract_frame_transforms(root_pos_w, root_quat_w, ee_pos_w, ee_quat_w)
# account for the offset
if self.cfg.body_offset is not None:
ee_pose_b, ee_quat_b = math_utils.combine_frame_transforms(
......@@ -395,7 +395,7 @@ class OperationalSpaceControllerAction(ActionTerm):
@property
def jacobian_b(self) -> torch.Tensor:
jacobian = self.jacobian_w
base_rot = self._asset.data.root_quat_w
base_rot = self._asset.data.root_link_quat_w
base_rot_matrix = math_utils.matrix_from_quat(math_utils.quat_inv(base_rot))
jacobian[:, :3, :] = torch.bmm(base_rot_matrix, jacobian[:, :3, :])
jacobian[:, 3:, :] = torch.bmm(base_rot_matrix, jacobian[:, 3:, :])
......@@ -556,10 +556,14 @@ 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[:] = self._asset.data.body_state_w[:, self._ee_body_idx, :7]
self._ee_pose_w[:, 0:3] = self._asset.data.body_link_pos_w[:, self._ee_body_idx]
self._ee_pose_w[:, 3:7] = self._asset.data.body_link_quat_w[:, self._ee_body_idx]
# Compute the pose of the ee body in the root frame
self._ee_pose_b_no_offset[:, 0:3], self._ee_pose_b_no_offset[:, 3:7] = math_utils.subtract_frame_transforms(
self._asset.data.root_pos_w, self._asset.data.root_quat_w, self._ee_pose_w[:, 0:3], self._ee_pose_w[:, 3:7]
self._asset.data.root_link_pos_w,
self._asset.data.root_link_quat_w,
self._ee_pose_w[:, 0:3],
self._ee_pose_w[:, 3:7],
)
# Account for the offset
if self.cfg.body_offset is not None:
......@@ -572,13 +576,17 @@ 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_vel_w[:, self._ee_body_idx, :]
self._ee_vel_w[:] = self._asset.data.body_com_vel_w[:, self._ee_body_idx, :]
# Compute the relative velocity in the world frame
relative_vel_w = self._ee_vel_w - self._asset.data.root_vel_w
relative_vel_w = self._ee_vel_w - self._asset.data.root_com_vel_w
# Convert ee velocities from world to root frame
self._ee_vel_b[:, 0:3] = math_utils.quat_rotate_inverse(self._asset.data.root_quat_w, relative_vel_w[:, 0:3])
self._ee_vel_b[:, 3:6] = math_utils.quat_rotate_inverse(self._asset.data.root_quat_w, relative_vel_w[:, 3:6])
self._ee_vel_b[:, 0:3] = math_utils.quat_rotate_inverse(
self._asset.data.root_link_quat_w, relative_vel_w[:, 0:3]
)
self._ee_vel_b[:, 3:6] = math_utils.quat_rotate_inverse(
self._asset.data.root_link_quat_w, relative_vel_w[:, 3:6]
)
# Account for the offset
if self.cfg.body_offset is not None:
......@@ -595,7 +603,7 @@ 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_quat_w, self._ee_force_w)
self._ee_force_b[:] = math_utils.quat_rotate_inverse(self._asset.data.root_link_quat_w, self._ee_force_w)
def _compute_task_frame_pose(self):
"""Computes the pose of the task frame in root frame."""
......@@ -604,8 +612,8 @@ 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_pos_w,
self._asset.data.root_quat_w,
self._asset.data.root_link_pos_w,
self._asset.data.root_link_quat_w,
self._task_frame_transformer.data.target_pos_w[:, 0, :],
self._task_frame_transformer.data.target_quat_w[:, 0, :],
)
......
......@@ -81,7 +81,9 @@ 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_pos_w[:, :2], dim=1)
self.metrics["error_pos_2d"] = torch.norm(
self.pos_command_w[:, :2] - self.robot.data.root_link_pos_w[:, :2], dim=1
)
self.metrics["error_heading"] = torch.abs(wrap_to_pi(self.heading_command_w - self.robot.data.heading_w))
def _resample_command(self, env_ids: Sequence[int]):
......@@ -95,7 +97,7 @@ 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_pos_w[env_ids]
target_vec = self.pos_command_w[env_ids] - self.robot.data.root_link_pos_w[env_ids]
target_direction = torch.atan2(target_vec[:, 1], target_vec[:, 0])
flipped_target_direction = wrap_to_pi(target_direction + torch.pi)
......@@ -116,8 +118,8 @@ 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_pos_w[:, :3]
self.pos_command_b[:] = quat_rotate_inverse(yaw_quat(self.robot.data.root_quat_w), target_vec)
target_vec = self.pos_command_w - self.robot.data.root_link_pos_w[:, :3]
self.pos_command_b[:] = quat_rotate_inverse(yaw_quat(self.robot.data.root_link_quat_w), target_vec)
self.heading_command_b[:] = wrap_to_pi(self.heading_command_w - self.robot.data.heading_w)
def _set_debug_vis_impl(self, debug_vis: bool):
......@@ -182,7 +184,7 @@ 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_pos_w[env_ids]
target_vec = self.pos_command_w[env_ids] - self.robot.data.root_link_pos_w[env_ids]
target_direction = torch.atan2(target_vec[:, 1], target_vec[:, 0])
flipped_target_direction = wrap_to_pi(target_direction + torch.pi)
......
......@@ -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_pos_w,
self.robot.data.root_quat_w,
self.robot.data.root_link_pos_w,
self.robot.data.root_link_quat_w,
self.pose_command_b[:, :3],
self.pose_command_b[:, 3:],
)
......@@ -101,8 +101,8 @@ class UniformPoseCommand(CommandTerm):
pos_error, rot_error = compute_pose_error(
self.pose_command_w[:, :3],
self.pose_command_w[:, 3:],
self.robot.data.body_state_w[:, self.body_idx, :3],
self.robot.data.body_state_w[:, self.body_idx, 3:7],
self.robot.data.body_link_state_w[:, self.body_idx, :3],
self.robot.data.body_link_state_w[:, self.body_idx, 3:7],
)
self.metrics["position_error"] = torch.norm(pos_error, dim=-1)
self.metrics["orientation_error"] = torch.norm(rot_error, dim=-1)
......@@ -151,5 +151,5 @@ class UniformPoseCommand(CommandTerm):
# -- goal pose
self.goal_pose_visualizer.visualize(self.pose_command_w[:, :3], self.pose_command_w[:, 3:])
# -- current body pose
body_pose_w = self.robot.data.body_state_w[:, self.body_idx]
self.current_pose_visualizer.visualize(body_pose_w[:, :3], body_pose_w[:, 3:7])
body_link_state_w = self.robot.data.body_link_state_w[:, self.body_idx]
self.current_pose_visualizer.visualize(body_link_state_w[:, :3], body_link_state_w[:, 3:7])
......@@ -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_lin_vel_b[:, :2], dim=-1) / max_command_step
torch.norm(self.vel_command_b[:, :2] - self.robot.data.root_com_lin_vel_b[:, :2], dim=-1) / max_command_step
)
self.metrics["error_vel_yaw"] += (
torch.abs(self.vel_command_b[:, 2] - self.robot.data.root_ang_vel_b[:, 2]) / max_command_step
torch.abs(self.vel_command_b[:, 2] - self.robot.data.root_com_ang_vel_b[:, 2]) / max_command_step
)
def _resample_command(self, env_ids: Sequence[int]):
......@@ -184,11 +184,11 @@ class UniformVelocityCommand(CommandTerm):
return
# get marker location
# -- base state
base_pos_w = self.robot.data.root_pos_w.clone()
base_pos_w = self.robot.data.root_link_pos_w.clone()
base_pos_w[:, 2] += 0.5
# -- resolve the scales and quaternions
vel_des_arrow_scale, vel_des_arrow_quat = self._resolve_xy_velocity_to_arrow(self.command[:, :2])
vel_arrow_scale, vel_arrow_quat = self._resolve_xy_velocity_to_arrow(self.robot.data.root_lin_vel_b[:, :2])
vel_arrow_scale, vel_arrow_quat = self._resolve_xy_velocity_to_arrow(self.robot.data.root_com_lin_vel_b[:, :2])
# display markers
self.goal_vel_visualizer.visualize(base_pos_w, vel_des_arrow_quat, vel_des_arrow_scale)
self.current_vel_visualizer.visualize(base_pos_w, vel_arrow_quat, vel_arrow_scale)
......@@ -209,7 +209,7 @@ 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_quat_w
base_quat_w = self.robot.data.root_link_quat_w
arrow_quat = math_utils.quat_mul(base_quat_w, arrow_quat)
return arrow_scale, arrow_quat
......
......@@ -624,13 +624,13 @@ def push_by_setting_velocity(
asset: RigidObject | Articulation = env.scene[asset_cfg.name]
# velocities
vel_w = asset.data.root_vel_w[env_ids]
vel_w = asset.data.root_com_vel_w[env_ids]
# sample random velocities
range_list = [velocity_range.get(key, (0.0, 0.0)) for key in ["x", "y", "z", "roll", "pitch", "yaw"]]
ranges = torch.tensor(range_list, device=asset.device)
vel_w[:] = math_utils.sample_uniform(ranges[:, 0], ranges[:, 1], vel_w.shape, device=asset.device)
# set the velocities into the physics simulation
asset.write_root_velocity_to_sim(vel_w, env_ids=env_ids)
asset.write_root_com_velocity_to_sim(vel_w, env_ids=env_ids)
def reset_root_state_uniform(
......@@ -674,8 +674,8 @@ def reset_root_state_uniform(
velocities = root_states[:, 7:13] + rand_samples
# set into the physics simulation
asset.write_root_pose_to_sim(torch.cat([positions, orientations], dim=-1), env_ids=env_ids)
asset.write_root_velocity_to_sim(velocities, env_ids=env_ids)
asset.write_root_link_pose_to_sim(torch.cat([positions, orientations], dim=-1), env_ids=env_ids)
asset.write_root_com_velocity_to_sim(velocities, env_ids=env_ids)
def reset_root_state_with_random_orientation(
......@@ -726,8 +726,8 @@ def reset_root_state_with_random_orientation(
velocities = root_states[:, 7:13] + rand_samples
# set into the physics simulation
asset.write_root_pose_to_sim(torch.cat([positions, orientations], dim=-1), env_ids=env_ids)
asset.write_root_velocity_to_sim(velocities, env_ids=env_ids)
asset.write_root_link_pose_to_sim(torch.cat([positions, orientations], dim=-1), env_ids=env_ids)
asset.write_root_com_velocity_to_sim(velocities, env_ids=env_ids)
def reset_root_state_from_terrain(
......@@ -793,8 +793,8 @@ def reset_root_state_from_terrain(
velocities = asset.data.default_root_state[:, 7:13] + rand_samples
# set into the physics simulation
asset.write_root_pose_to_sim(torch.cat([positions, orientations], dim=-1), env_ids=env_ids)
asset.write_root_velocity_to_sim(velocities, env_ids=env_ids)
asset.write_root_link_pose_to_sim(torch.cat([positions, orientations], dim=-1), env_ids=env_ids)
asset.write_root_com_velocity_to_sim(velocities, env_ids=env_ids)
def reset_joints_by_scale(
......@@ -914,14 +914,16 @@ def reset_scene_to_default(env: ManagerBasedEnv, env_ids: torch.Tensor):
default_root_state = rigid_object.data.default_root_state[env_ids].clone()
default_root_state[:, 0:3] += env.scene.env_origins[env_ids]
# set into the physics simulation
rigid_object.write_root_state_to_sim(default_root_state, env_ids=env_ids)
rigid_object.write_root_link_pose_to_sim(default_root_state[:, :7], env_ids=env_ids)
rigid_object.write_root_com_velocity_to_sim(default_root_state[:, 7:], env_ids=env_ids)
# articulations
for articulation_asset in env.scene.articulations.values():
# obtain default and deal with the offset for env origins
default_root_state = articulation_asset.data.default_root_state[env_ids].clone()
default_root_state[:, 0:3] += env.scene.env_origins[env_ids]
# set into the physics simulation
articulation_asset.write_root_state_to_sim(default_root_state, env_ids=env_ids)
articulation_asset.write_root_link_pose_to_sim(default_root_state[:, :7], env_ids=env_ids)
articulation_asset.write_root_com_velocity_to_sim(default_root_state[:, 7:], env_ids=env_ids)
# obtain default joint positions
default_joint_pos = articulation_asset.data.default_joint_pos[env_ids].clone()
default_joint_vel = articulation_asset.data.default_joint_vel[env_ids].clone()
......
......@@ -34,21 +34,21 @@ def base_pos_z(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg(
"""Root height in the simulation world frame."""
# extract the used quantities (to enable type-hinting)
asset: Articulation = env.scene[asset_cfg.name]
return asset.data.root_pos_w[:, 2].unsqueeze(-1)
return asset.data.root_link_pos_w[:, 2].unsqueeze(-1)
def base_lin_vel(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor:
"""Root linear velocity in the asset's root frame."""
# extract the used quantities (to enable type-hinting)
asset: RigidObject = env.scene[asset_cfg.name]
return asset.data.root_lin_vel_b
return asset.data.root_com_lin_vel_b
def base_ang_vel(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor:
"""Root angular velocity in the asset's root frame."""
# extract the used quantities (to enable type-hinting)
asset: RigidObject = env.scene[asset_cfg.name]
return asset.data.root_ang_vel_b
return asset.data.root_com_ang_vel_b
def projected_gravity(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor:
......@@ -62,7 +62,7 @@ def root_pos_w(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg(
"""Asset root position in the environment frame."""
# extract the used quantities (to enable type-hinting)
asset: RigidObject = env.scene[asset_cfg.name]
return asset.data.root_pos_w - env.scene.env_origins
return asset.data.root_link_pos_w - env.scene.env_origins
def root_quat_w(
......@@ -77,7 +77,7 @@ def root_quat_w(
# extract the used quantities (to enable type-hinting)
asset: RigidObject = env.scene[asset_cfg.name]
quat = asset.data.root_quat_w
quat = asset.data.root_link_quat_w
# make the quaternion real-part positive if configured
return math_utils.quat_unique(quat) if make_quat_unique else quat
......@@ -86,14 +86,14 @@ def root_lin_vel_w(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntity
"""Asset root linear velocity in the environment frame."""
# extract the used quantities (to enable type-hinting)
asset: RigidObject = env.scene[asset_cfg.name]
return asset.data.root_lin_vel_w
return asset.data.root_com_lin_vel_w
def root_ang_vel_w(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor:
"""Asset root angular velocity in the environment frame."""
# extract the used quantities (to enable type-hinting)
asset: RigidObject = env.scene[asset_cfg.name]
return asset.data.root_ang_vel_w
return asset.data.root_com_ang_vel_w
"""
......
......@@ -77,14 +77,14 @@ def lin_vel_z_l2(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg = SceneEntity
"""Penalize z-axis base linear velocity using L2 squared kernel."""
# extract the used quantities (to enable type-hinting)
asset: RigidObject = env.scene[asset_cfg.name]
return torch.square(asset.data.root_lin_vel_b[:, 2])
return torch.square(asset.data.root_com_lin_vel_b[:, 2])
def ang_vel_xy_l2(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor:
"""Penalize xy-axis base angular velocity using L2 squared kernel."""
# extract the used quantities (to enable type-hinting)
asset: RigidObject = env.scene[asset_cfg.name]
return torch.sum(torch.square(asset.data.root_ang_vel_b[:, :2]), dim=1)
return torch.sum(torch.square(asset.data.root_com_ang_vel_b[:, :2]), dim=1)
def flat_orientation_l2(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor:
......@@ -119,7 +119,7 @@ def base_height_l2(
# Use the provided target height directly for flat terrain
adjusted_target_height = target_height
# Compute the L2 squared penalty
return torch.square(asset.data.root_pos_w[:, 2] - adjusted_target_height)
return torch.square(asset.data.root_link_pos_w[:, 2] - adjusted_target_height)
def body_lin_acc_l2(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor:
......@@ -292,7 +292,7 @@ def track_lin_vel_xy_exp(
asset: RigidObject = env.scene[asset_cfg.name]
# compute the error
lin_vel_error = torch.sum(
torch.square(env.command_manager.get_command(command_name)[:, :2] - asset.data.root_lin_vel_b[:, :2]),
torch.square(env.command_manager.get_command(command_name)[:, :2] - asset.data.root_com_lin_vel_b[:, :2]),
dim=1,
)
return torch.exp(-lin_vel_error / std**2)
......@@ -305,5 +305,7 @@ def track_ang_vel_z_exp(
# extract the used quantities (to enable type-hinting)
asset: RigidObject = env.scene[asset_cfg.name]
# compute the error
ang_vel_error = torch.square(env.command_manager.get_command(command_name)[:, 2] - asset.data.root_ang_vel_b[:, 2])
ang_vel_error = torch.square(
env.command_manager.get_command(command_name)[:, 2] - asset.data.root_com_ang_vel_b[:, 2]
)
return torch.exp(-ang_vel_error / std**2)
......@@ -69,7 +69,7 @@ def root_height_below_minimum(
"""
# extract the used quantities (to enable type-hinting)
asset: RigidObject = env.scene[asset_cfg.name]
return asset.data.root_pos_w[:, 2] < minimum_height
return asset.data.root_link_pos_w[:, 2] < minimum_height
"""
......
......@@ -156,7 +156,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_pos_w[self.cfg.env_index]
self.viewer_origin = self._env.scene[self.cfg.asset_name].data.root_link_pos_w[self.cfg.env_index]
# update the camera view
self.update_view_location()
......
......@@ -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_state_w[:, :7].clone()
asset_state["root_pose"] = articulation.data.root_link_state_w[:, :7].clone()
if is_relative:
asset_state["root_pose"][:, :3] -= self.env_origins
asset_state["root_velocity"] = articulation.data.root_vel_w.clone()
asset_state["root_velocity"] = articulation.data.root_com_vel_w.clone()
asset_state["joint_position"] = articulation.data.joint_pos.clone()
asset_state["joint_velocity"] = articulation.data.joint_vel.clone()
state["articulation"][asset_name] = asset_state
......@@ -384,10 +384,10 @@ class InteractiveScene:
state["rigid_object"] = dict()
for asset_name, rigid_object in self._rigid_objects.items():
asset_state = dict()
asset_state["root_pose"] = rigid_object.data.root_state_w[:, :7].clone()
asset_state["root_pose"] = rigid_object.data.root_link_state_w[:, :7].clone()
if is_relative:
asset_state["root_pose"][:, :3] -= self.env_origins
asset_state["root_velocity"] = rigid_object.data.root_vel_w.clone()
asset_state["root_velocity"] = rigid_object.data.root_com_vel_w.clone()
state["rigid_object"][asset_name] = asset_state
return state
......@@ -439,8 +439,8 @@ class InteractiveScene:
if is_relative:
root_pose[:, :3] += self.env_origins[env_ids]
root_velocity = asset_state["root_velocity"].clone()
articulation.write_root_pose_to_sim(root_pose, env_ids=env_ids)
articulation.write_root_velocity_to_sim(root_velocity, env_ids=env_ids)
articulation.write_root_link_pose_to_sim(root_pose, env_ids=env_ids)
articulation.write_root_com_velocity_to_sim(root_velocity, env_ids=env_ids)
# joint state
joint_position = asset_state["joint_position"].clone()
joint_velocity = asset_state["joint_velocity"].clone()
......@@ -463,8 +463,8 @@ class InteractiveScene:
if is_relative:
root_pose[:, :3] += self.env_origins[env_ids]
root_velocity = asset_state["root_velocity"].clone()
rigid_object.write_root_pose_to_sim(root_pose, env_ids=env_ids)
rigid_object.write_root_velocity_to_sim(root_velocity, env_ids=env_ids)
rigid_object.write_root_link_pose_to_sim(root_pose, env_ids=env_ids)
rigid_object.write_root_com_velocity_to_sim(root_velocity, env_ids=env_ids)
self.write_data_to_sim()
def write_data_to_sim(self):
......
......@@ -746,7 +746,7 @@ def is_identity_pose(pos: torch.tensor, rot: torch.tensor) -> bool:
return torch.allclose(pos, pos_identity) and torch.allclose(rot, rot_identity)
# @torch.jit.script
@torch.jit.script
def combine_frame_transforms(
t01: torch.Tensor, q01: torch.Tensor, t12: torch.Tensor | None = None, q12: torch.Tensor | None = None
) -> tuple[torch.Tensor, torch.Tensor]:
......
......@@ -99,7 +99,8 @@ def main():
root_state = robot.data.default_root_state.clone()
root_state[0, :2] = torch.tensor([0.0, -0.5], device=sim.device)
root_state[1, :2] = torch.tensor([0.0, 0.5], device=sim.device)
robot.write_root_state_to_sim(root_state)
robot.write_root_link_pose_to_sim(root_state[:, :7])
robot.write_root_com_velocity_to_sim(root_state[:, 7:])
# reset dof state
joint_pos, joint_vel = robot.data.default_joint_pos, robot.data.default_joint_vel
robot.write_joint_state_to_sim(joint_pos, joint_vel)
......
......@@ -113,7 +113,8 @@ def run_simulator(sim: sim_utils.SimulationContext, entities: dict[str, Articula
root_state = robot.data.default_root_state.clone()
root_state[:, :3] += origins[index]
root_state[:, :2] += torch.randn_like(root_state[:, :2]) * 0.25
robot.write_root_state_to_sim(root_state)
robot.write_root_link_pose_to_sim(root_state[:, :7])
robot.write_root_com_velocity_to_sim(root_state[:, 7:])
# joint state
joint_pos, joint_vel = robot.data.default_joint_pos.clone(), robot.data.default_joint_vel.clone()
robot.write_joint_state_to_sim(joint_pos, joint_vel)
......
......@@ -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_state_w[:, ee_frame_idx, 0:7]
root_pose_w = robot.data.root_state_w[:, 0:7]
ee_pose_w = robot.data.body_link_state_w[:, ee_frame_idx, 0:7]
root_pose_w = robot.data.root_link_state_w[:, 0:7]
ee_pos_b, ee_quat_b = subtract_frame_transforms(
root_pose_w[:, 0:3], root_pose_w[:, 3:7], ee_pose_w[:, 0:3], ee_pose_w[:, 3:7]
)
......@@ -177,9 +177,10 @@ 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_state_w.clone()
root_state = robot.data.root_link_state_w.clone()
root_state[:, 3:7] = random_yaw_orientation(self.num_envs, self.sim.device)
robot.write_root_state_to_sim(root_state)
robot.write_root_link_pose_to_sim(root_state[:, :7])
robot.write_root_com_velocity_to_sim(root_state[:, 7:])
robot.reset()
# reset actions
ee_pose_b_des[:] = self.ee_pose_b_des_set[current_goal_idx]
......@@ -194,8 +195,8 @@ 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_state_w[:, ee_frame_idx, 0:7]
root_pose_w = robot.data.root_state_w[:, 0:7]
ee_pose_w = robot.data.body_link_state_w[:, ee_frame_idx, 0:7]
root_pose_w = robot.data.root_link_state_w[:, 0:7]
base_rot = root_pose_w[:, 3:7]
base_rot_matrix = matrix_from_quat(quat_inv(base_rot))
jacobian[:, :3, :] = torch.bmm(base_rot_matrix, jacobian[:, :3, :])
......
......@@ -723,24 +723,26 @@ 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_state_w[:, 3:7]))
root_rot_matrix = matrix_from_quat(quat_inv(robot.data.root_link_quat_w))
jacobian_b[:, :3, :] = torch.bmm(root_rot_matrix, jacobian_b[:, :3, :])
jacobian_b[:, 3:, :] = torch.bmm(root_rot_matrix, jacobian_b[:, 3:, :])
# Compute current pose of the end-effector
root_pose_w = robot.data.root_state_w[:, 0:7]
ee_pose_w = robot.data.body_state_w[:, ee_frame_idx, 0:7]
root_pose_w = robot.data.root_link_state_w[:, 0:7]
ee_pose_w = robot.data.body_link_state_w[:, ee_frame_idx, 0:7]
ee_pos_b, ee_quat_b = subtract_frame_transforms(
root_pose_w[:, 0:3], root_pose_w[:, 3:7], ee_pose_w[:, 0:3], ee_pose_w[:, 3:7]
)
ee_pose_b = torch.cat([ee_pos_b, ee_quat_b], dim=-1)
# Compute the current velocity of the end-effector
ee_vel_w = robot.data.body_vel_w[:, ee_frame_idx, :] # Extract end-effector velocity in the world frame
root_vel_w = robot.data.root_vel_w # Extract root velocity in the world frame
ee_vel_w = robot.data.body_com_vel_w[:, ee_frame_idx, :] # Extract end-effector velocity in the world frame
root_vel_w = robot.data.root_com_vel_w # Extract root velocity in the world frame
relative_vel_w = ee_vel_w - root_vel_w # Compute the relative velocity in the world frame
ee_lin_vel_b = quat_rotate_inverse(robot.data.root_quat_w, relative_vel_w[:, 0:3]) # From world to root frame
ee_ang_vel_b = quat_rotate_inverse(robot.data.root_quat_w, relative_vel_w[:, 3:6])
ee_lin_vel_b = quat_rotate_inverse(
robot.data.root_link_quat_w, relative_vel_w[:, 0:3]
) # From world to root frame
ee_ang_vel_b = quat_rotate_inverse(robot.data.root_link_quat_w, relative_vel_w[:, 3:6])
ee_vel_b = torch.cat([ee_lin_vel_b, ee_ang_vel_b], dim=-1)
# Calculate the contact force
......
......@@ -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_pos_w - self._env.scene.env_origins)
vel_error = -self._asset.data.root_lin_vel_w
pos_error = self._processed_actions - (self._asset.data.root_link_pos_w - self._env.scene.env_origins)
vel_error = -self._asset.data.root_com_lin_vel_w
# set velocity targets
self._vel_command[:, :3] = self.p_gain * pos_error + self.d_gain * vel_error
self._asset.write_root_velocity_to_sim(self._vel_command)
self._asset.write_root_com_velocity_to_sim(self._vel_command)
@configclass
......@@ -151,7 +151,7 @@ def base_position(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg) -> torch.Tens
"""Root linear velocity in the asset's root frame."""
# extract the used quantities (to enable type-hinting)
asset: RigidObject = env.scene[asset_cfg.name]
return asset.data.root_pos_w - env.scene.env_origins
return asset.data.root_link_pos_w - env.scene.env_origins
##
......
......@@ -98,7 +98,8 @@ def run_simulator(
# root state
root_state = scene["robot"].data.default_root_state.clone()
root_state[:, :3] += scene.env_origins
scene["robot"].write_root_state_to_sim(root_state)
scene["robot"].write_root_link_pose_to_sim(root_state[:, :7])
scene["robot"].write_root_com_velocity_to_sim(root_state[:, 7:])
# set joint positions with some noise
joint_pos, joint_vel = (
scene["robot"].data.default_joint_pos.clone(),
......
......@@ -130,11 +130,13 @@ def main():
joint_vel = scene.articulations["robot_1"].data.default_joint_vel
# -- set root state
# -- robot 1
scene.articulations["robot_1"].write_root_state_to_sim(root_state)
scene.articulations["robot_1"].write_root_link_pose_to_sim(root_state[:, :7])
scene.articulations["robot_1"].write_root_com_velocity_to_sim(root_state[:, 7:])
scene.articulations["robot_1"].write_joint_state_to_sim(joint_pos, joint_vel)
# -- robot 2
root_state[:, 1] += 1.0
scene.articulations["robot_2"].write_root_state_to_sim(root_state)
scene.articulations["robot_2"].write_root_link_pose_to_sim(root_state[:, :7])
scene.articulations["robot_2"].write_root_com_velocity_to_sim(root_state[:, 7:])
scene.articulations["robot_2"].write_joint_state_to_sim(joint_pos, joint_vel)
# reset buffers
scene.reset()
......
......@@ -151,8 +151,8 @@ def main():
# Get the ball initial positions
sim.step(render=not args_cli.headless)
balls.update(sim.get_physics_dt())
ball_initial_positions = balls.data.root_pos_w.clone()
ball_initial_orientations = balls.data.root_quat_w.clone()
ball_initial_positions = balls.data.root_link_pos_w.clone()
ball_initial_orientations = balls.data.root_link_quat_w.clone()
# Create a counter for resetting the scene
step_count = 0
......@@ -168,7 +168,7 @@ def main():
# Reset the scene
if step_count % 500 == 0:
# reset ball positions
balls.write_root_pose_to_sim(torch.cat([ball_initial_positions, ball_initial_orientations], dim=-1))
balls.write_root_link_pose_to_sim(torch.cat([ball_initial_positions, ball_initial_orientations], dim=-1))
balls.reset()
# reset the sensor
imu.reset()
......
......@@ -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_pose_to_sim(root_pose=test_pose)
shape.write_root_link_pose_to_sim(root_pose=test_pose)
# perform simulation step
self._perform_sim_step()
# increment contact time
......@@ -432,7 +432,7 @@ 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_pose_to_sim(root_pose=reset_pose)
shape.write_root_link_pose_to_sim(root_pose=reset_pose)
# perform simulation step
self._perform_sim_step()
# set the last air time to 2 sim_dt steps, because last_air_time and last_contact_time
......
......@@ -175,7 +175,8 @@ class TestFrameTransformer(unittest.TestCase):
joint_vel = scene.articulations["robot"].data.default_joint_vel
# -- set root state
# -- robot
scene.articulations["robot"].write_root_state_to_sim(root_state)
scene.articulations["robot"].write_root_link_pose_to_sim(root_state[:, :7])
scene.articulations["robot"].write_root_com_velocity_to_sim(root_state[:, 7:])
scene.articulations["robot"].write_joint_state_to_sim(joint_pos, joint_vel)
# reset buffers
scene.reset()
......@@ -192,9 +193,9 @@ class TestFrameTransformer(unittest.TestCase):
# check absolute frame transforms in world frame
# -- ground-truth
root_pose_w = scene.articulations["robot"].data.root_state_w[:, :7]
feet_pos_w_gt = scene.articulations["robot"].data.body_pos_w[:, feet_indices]
feet_quat_w_gt = scene.articulations["robot"].data.body_quat_w[:, feet_indices]
root_pose_w = scene.articulations["robot"].data.root_link_state_w[:, :7]
feet_pos_w_gt = scene.articulations["robot"].data.body_link_pos_w[:, feet_indices]
feet_quat_w_gt = scene.articulations["robot"].data.body_link_quat_w[:, feet_indices]
# -- frame transformer
source_pos_w_tf = scene.sensors["frame_transformer"].data.source_pos_w
source_quat_w_tf = scene.sensors["frame_transformer"].data.source_quat_w
......@@ -275,7 +276,8 @@ class TestFrameTransformer(unittest.TestCase):
joint_vel = scene.articulations["robot"].data.default_joint_vel
# -- set root state
# -- robot
scene.articulations["robot"].write_root_state_to_sim(root_state)
scene.articulations["robot"].write_root_link_pose_to_sim(root_state[:, :7])
scene.articulations["robot"].write_root_com_velocity_to_sim(root_state[:, 7:])
scene.articulations["robot"].write_joint_state_to_sim(joint_pos, joint_vel)
# reset buffers
scene.reset()
......@@ -292,9 +294,9 @@ class TestFrameTransformer(unittest.TestCase):
# check absolute frame transforms in world frame
# -- ground-truth
source_pose_w_gt = scene.articulations["robot"].data.body_state_w[:, source_frame_index, :7]
feet_pos_w_gt = scene.articulations["robot"].data.body_pos_w[:, feet_indices]
feet_quat_w_gt = scene.articulations["robot"].data.body_quat_w[:, feet_indices]
source_pose_w_gt = scene.articulations["robot"].data.body_link_state_w[:, source_frame_index, :7]
feet_pos_w_gt = scene.articulations["robot"].data.body_link_pos_w[:, feet_indices]
feet_quat_w_gt = scene.articulations["robot"].data.body_link_quat_w[:, feet_indices]
# -- frame transformer
source_pos_w_tf = scene.sensors["frame_transformer"].data.source_pos_w
source_quat_w_tf = scene.sensors["frame_transformer"].data.source_quat_w
......@@ -356,7 +358,8 @@ class TestFrameTransformer(unittest.TestCase):
joint_vel = scene.articulations["robot"].data.default_joint_vel
# -- set root state
# -- robot
scene.articulations["robot"].write_root_state_to_sim(root_state)
scene.articulations["robot"].write_root_link_pose_to_sim(root_state[:, :7])
scene.articulations["robot"].write_root_com_velocity_to_sim(root_state[:, 7:])
scene.articulations["robot"].write_joint_state_to_sim(joint_pos, joint_vel)
# reset buffers
scene.reset()
......@@ -373,9 +376,9 @@ class TestFrameTransformer(unittest.TestCase):
# check absolute frame transforms in world frame
# -- ground-truth
root_pose_w = scene.articulations["robot"].data.root_state_w[:, :7]
cube_pos_w_gt = scene.rigid_objects["cube"].data.root_state_w[:, :3]
cube_quat_w_gt = scene.rigid_objects["cube"].data.root_state_w[:, 3:7]
root_pose_w = scene.articulations["robot"].data.root_link_state_w[:, :7]
cube_pos_w_gt = scene.rigid_objects["cube"].data.root_link_state_w[:, :3]
cube_quat_w_gt = scene.rigid_objects["cube"].data.root_link_state_w[:, 3:7]
# -- frame transformer
source_pos_w_tf = scene.sensors["frame_transformer"].data.source_pos_w
source_quat_w_tf = scene.sensors["frame_transformer"].data.source_quat_w
......@@ -447,7 +450,8 @@ class TestFrameTransformer(unittest.TestCase):
root_state[:, :3] += scene.env_origins
# -- set root state
# -- cube
scene["cube"].write_root_state_to_sim(root_state)
scene["cube"].write_root_link_pose_to_sim(root_state[:, :7])
scene["cube"].write_root_com_velocity_to_sim(root_state[:, 7:])
# reset buffers
scene.reset()
......@@ -460,8 +464,8 @@ class TestFrameTransformer(unittest.TestCase):
# check absolute frame transforms in world frame
# -- ground-truth
cube_pos_w_gt = scene["cube"].data.root_state_w[:, :3]
cube_quat_w_gt = scene["cube"].data.root_state_w[:, 3:7]
cube_pos_w_gt = scene["cube"].data.root_link_state_w[:, :3]
cube_quat_w_gt = scene["cube"].data.root_link_state_w[:, 3:7]
# -- frame transformer
source_pos_w_tf = scene.sensors["frame_transformer"].data.source_pos_w
source_quat_w_tf = scene.sensors["frame_transformer"].data.source_quat_w
......@@ -534,7 +538,8 @@ class TestFrameTransformer(unittest.TestCase):
joint_vel = scene.articulations["robot"].data.default_joint_vel
# -- set root state
# -- robot
scene.articulations["robot"].write_root_state_to_sim(root_state)
scene.articulations["robot"].write_root_link_pose_to_sim(root_state[:, :7])
scene.articulations["robot"].write_root_com_velocity_to_sim(root_state[:, 7:])
scene.articulations["robot"].write_joint_state_to_sim(joint_pos, joint_vel)
# reset buffers
scene.reset()
......@@ -551,9 +556,9 @@ class TestFrameTransformer(unittest.TestCase):
# check absolute frame transforms in world frame
# -- ground-truth
root_pose_w = scene.articulations["robot"].data.root_state_w[:, :7]
bodies_pos_w_gt = scene.articulations["robot"].data.body_pos_w
bodies_quat_w_gt = scene.articulations["robot"].data.body_quat_w
root_pose_w = scene.articulations["robot"].data.root_link_state_w[:, :7]
bodies_pos_w_gt = scene.articulations["robot"].data.body_link_pos_w
bodies_quat_w_gt = scene.articulations["robot"].data.body_link_quat_w
# -- frame transformer
source_pos_w_tf = scene.sensors["frame_transformer"].data.source_pos_w
......
......@@ -189,12 +189,12 @@ class TestImu(unittest.TestCase):
for idx in range(200):
# set velocity
self.scene.rigid_objects["balls"].write_root_velocity_to_sim(
self.scene.rigid_objects["balls"].write_root_com_velocity_to_sim(
torch.tensor([[1.0, 0.0, 0.0, 0.0, 0.0, 0.0]], dtype=torch.float32, device=self.scene.device).repeat(
self.scene.num_envs, 1
)
)
self.scene.rigid_objects["cube"].write_root_velocity_to_sim(
self.scene.rigid_objects["cube"].write_root_com_velocity_to_sim(
torch.tensor([[1.0, 0.0, 0.0, 0.0, 0.0, 0.0]], dtype=torch.float32, device=self.scene.device).repeat(
self.scene.num_envs, 1
)
......@@ -236,7 +236,7 @@ class TestImu(unittest.TestCase):
)
# check the imu velocities
# NOTE: the expected lin_vel_b is the same as the set velocity, as write_root_velocity_to_sim is
# NOTE: the expected lin_vel_b is the same as the set velocity, as write_root_com_velocity_to_sim is
# setting v_0 (initial velocity) and then a calculation step of v_i = v_0 + a*dt. Consequently,
# the data.lin_vel_b is returning approx. v_i.
torch.testing.assert_close(
......@@ -266,7 +266,7 @@ 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_velocity_to_sim(
self.scene.rigid_objects["balls"].write_root_com_velocity_to_sim(
torch.tensor([[0.1, 0.0, 0.0, 0.0, 0.0, 0.0]], dtype=torch.float32, device=self.scene.device).repeat(
self.scene.num_envs, 1
)
......@@ -287,7 +287,7 @@ 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_quat_w,
self.scene.rigid_objects["balls"].data.root_link_quat_w,
torch.tensor([[0.1, 0.0, 0.0]], dtype=torch.float32, device=self.scene.device).repeat(
self.scene.num_envs, 1
)
......@@ -300,7 +300,7 @@ 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_ang_vel_b,
self.scene.rigid_objects["balls"].data.root_com_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_velocity_to_sim(
self.scene.articulations["robot"].write_root_com_velocity_to_sim(
torch.tensor([[0.05, 0.0, 0.0, 0.0, 0.0, 0.0]], dtype=torch.float32, device=self.scene.device).repeat(
self.scene.num_envs, 1
)
......@@ -504,7 +504,7 @@ class TestImu(unittest.TestCase):
for idx in range(10):
# set acceleration
self.scene.articulations["robot"].write_root_velocity_to_sim(
self.scene.articulations["robot"].write_root_com_velocity_to_sim(
torch.tensor([[0.5, 0.0, 0.0, 0.0, 0.0, 0.0]], dtype=torch.float32, device=self.scene.device).repeat(
self.scene.num_envs, 1
)
......
......@@ -89,8 +89,8 @@ class AnymalCEnv(DirectRLEnv):
[
tensor
for tensor in (
self._robot.data.root_lin_vel_b,
self._robot.data.root_ang_vel_b,
self._robot.data.root_com_lin_vel_b,
self._robot.data.root_com_ang_vel_b,
self._robot.data.projected_gravity_b,
self._commands,
self._robot.data.joint_pos - self._robot.data.default_joint_pos,
......@@ -107,15 +107,17 @@ 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_lin_vel_b[:, :2]), dim=1)
lin_vel_error = torch.sum(
torch.square(self._commands[:, :2] - self._robot.data.root_com_lin_vel_b[:, :2]), dim=1
)
lin_vel_error_mapped = torch.exp(-lin_vel_error / 0.25)
# yaw rate tracking
yaw_rate_error = torch.square(self._commands[:, 2] - self._robot.data.root_ang_vel_b[:, 2])
yaw_rate_error = torch.square(self._commands[:, 2] - self._robot.data.root_com_ang_vel_b[:, 2])
yaw_rate_error_mapped = torch.exp(-yaw_rate_error / 0.25)
# z velocity tracking
z_vel_error = torch.square(self._robot.data.root_lin_vel_b[:, 2])
z_vel_error = torch.square(self._robot.data.root_com_lin_vel_b[:, 2])
# angular velocity x/y
ang_vel_error = torch.sum(torch.square(self._robot.data.root_ang_vel_b[:, :2]), dim=1)
ang_vel_error = torch.sum(torch.square(self._robot.data.root_com_ang_vel_b[:, :2]), dim=1)
# joint torques
joint_torques = torch.sum(torch.square(self._robot.data.applied_torque), dim=1)
# joint acceleration
......@@ -178,8 +180,8 @@ 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_pose_to_sim(default_root_state[:, :7], env_ids)
self._robot.write_root_velocity_to_sim(default_root_state[:, 7:], env_ids)
self._robot.write_root_link_pose_to_sim(default_root_state[:, :7], env_ids)
self._robot.write_root_com_velocity_to_sim(default_root_state[:, 7:], env_ids)
self._robot.write_joint_state_to_sim(joint_pos, joint_vel, None, env_ids)
# Logging
extras = dict()
......
......@@ -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_pose_to_sim(default_root_state[:, :7], env_ids)
self.robot.write_root_velocity_to_sim(default_root_state[:, 7:], env_ids)
self.robot.write_root_link_pose_to_sim(default_root_state[:, :7], env_ids)
self.robot.write_root_com_velocity_to_sim(default_root_state[:, 7:], env_ids)
self.robot.write_joint_state_to_sim(joint_pos, joint_vel, None, env_ids)
......
......@@ -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_pose_to_sim(default_root_state[:, :7], env_ids)
self._cartpole.write_root_velocity_to_sim(default_root_state[:, 7:], env_ids)
self._cartpole.write_root_link_pose_to_sim(default_root_state[:, :7], env_ids)
self._cartpole.write_root_com_velocity_to_sim(default_root_state[:, 7:], env_ids)
self._cartpole.write_joint_state_to_sim(joint_pos, joint_vel, None, env_ids)
......
......@@ -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_pose_to_sim(default_root_state[:, :7], env_ids)
self.cartpole.write_root_velocity_to_sim(default_root_state[:, 7:], env_ids)
self.cartpole.write_root_link_pose_to_sim(default_root_state[:, :7], env_ids)
self.cartpole.write_root_com_velocity_to_sim(default_root_state[:, 7:], env_ids)
self.cartpole.write_joint_state_to_sim(joint_pos, joint_vel, None, env_ids)
......
......@@ -189,16 +189,18 @@ 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_pos_w - self.scene.env_origins
self.fixed_quat = self._fixed_asset.data.root_quat_w
self.fixed_pos = self._fixed_asset.data.root_link_pos_w - self.scene.env_origins
self.fixed_quat = self._fixed_asset.data.root_link_quat_w
self.held_pos = self._held_asset.data.root_pos_w - self.scene.env_origins
self.held_quat = self._held_asset.data.root_quat_w
self.held_pos = self._held_asset.data.root_link_pos_w - self.scene.env_origins
self.held_quat = self._held_asset.data.root_link_quat_w
self.fingertip_midpoint_pos = self._robot.data.body_pos_w[:, self.fingertip_body_idx] - self.scene.env_origins
self.fingertip_midpoint_quat = self._robot.data.body_quat_w[:, self.fingertip_body_idx]
self.fingertip_midpoint_linvel = self._robot.data.body_lin_vel_w[:, self.fingertip_body_idx]
self.fingertip_midpoint_angvel = self._robot.data.body_ang_vel_w[:, self.fingertip_body_idx]
self.fingertip_midpoint_pos = (
self._robot.data.body_link_pos_w[:, self.fingertip_body_idx] - self.scene.env_origins
)
self.fingertip_midpoint_quat = self._robot.data.body_link_quat_w[:, self.fingertip_body_idx]
self.fingertip_midpoint_linvel = self._robot.data.body_com_lin_vel_w[:, self.fingertip_body_idx]
self.fingertip_midpoint_angvel = self._robot.data.body_com_ang_vel_w[:, self.fingertip_body_idx]
jacobians = self._robot.root_physx_view.get_jacobians()
......@@ -552,13 +554,15 @@ 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_state_to_sim(held_state, env_ids=env_ids)
self._held_asset.write_root_link_pose_to_sim(held_state[:, 0:7], env_ids=env_ids)
self._held_asset.write_root_com_velocity_to_sim(held_state[:, 7:], env_ids=env_ids)
self._held_asset.reset()
fixed_state = self._fixed_asset.data.default_root_state.clone()[env_ids]
fixed_state[:, 0:3] += self.scene.env_origins[env_ids]
fixed_state[:, 7:] = 0.0
self._fixed_asset.write_root_state_to_sim(fixed_state, env_ids=env_ids)
self._fixed_asset.write_root_link_pose_to_sim(fixed_state[:, 0:7], env_ids=env_ids)
self._fixed_asset.write_root_com_velocity_to_sim(fixed_state[:, 7:], env_ids=env_ids)
self._fixed_asset.reset()
def set_pos_inverse_kinematics(self, env_ids):
......@@ -681,7 +685,8 @@ class FactoryEnv(DirectRLEnv):
# (1.c.) Velocity
fixed_state[:, 7:] = 0.0 # vel
# (1.d.) Update values.
self._fixed_asset.write_root_state_to_sim(fixed_state, env_ids=env_ids)
self._fixed_asset.write_root_link_pose_to_sim(fixed_state[:, 0:7], env_ids=env_ids)
self._fixed_asset.write_root_com_velocity_to_sim(fixed_state[:, 7:], env_ids=env_ids)
self._fixed_asset.reset()
# (1.e.) Noisy position observation.
......@@ -767,13 +772,15 @@ 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_state_to_sim(small_gear_state, env_ids=env_ids)
self._small_gear_asset.write_root_link_pose_to_sim(small_gear_state[:, 0:7], env_ids=env_ids)
self._small_gear_asset.write_root_com_velocity_to_sim(small_gear_state[:, 7:], env_ids=env_ids)
self._small_gear_asset.reset()
large_gear_state = self._large_gear_asset.data.default_root_state.clone()[env_ids]
large_gear_state[:, 0:7] = fixed_state[:, 0:7]
large_gear_state[:, 7:] = 0.0 # vel
self._large_gear_asset.write_root_state_to_sim(large_gear_state, env_ids=env_ids)
self._large_gear_asset.write_root_link_pose_to_sim(large_gear_state[:, 0:7], env_ids=env_ids)
self._large_gear_asset.write_root_com_velocity_to_sim(large_gear_state[:, 7:], env_ids=env_ids)
self._large_gear_asset.reset()
# (3) Randomize asset-in-gripper location.
......@@ -815,7 +822,8 @@ 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_state_to_sim(held_state)
self._held_asset.write_root_link_pose_to_sim(held_state[:, 0:7])
self._held_asset.write_root_com_velocity_to_sim(held_state[:, 7:])
self._held_asset.reset()
# Close hand
......
......@@ -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_pos_w[:, self.left_finger_link_idx]
robot_right_finger_pos = self._robot.data.body_pos_w[:, self.right_finger_link_idx]
robot_left_finger_pos = self._robot.data.body_link_pos_w[:, self.left_finger_link_idx]
robot_right_finger_pos = self._robot.data.body_link_pos_w[:, self.right_finger_link_idx]
return self._compute_rewards(
self.actions,
......@@ -372,10 +372,10 @@ class FrankaCabinetEnv(DirectRLEnv):
if env_ids is None:
env_ids = self._robot._ALL_INDICES
hand_pos = self._robot.data.body_pos_w[env_ids, self.hand_link_idx]
hand_rot = self._robot.data.body_quat_w[env_ids, self.hand_link_idx]
drawer_pos = self._cabinet.data.body_pos_w[env_ids, self.drawer_link_idx]
drawer_rot = self._cabinet.data.body_quat_w[env_ids, self.drawer_link_idx]
hand_pos = self._robot.data.body_link_pos_w[env_ids, self.hand_link_idx]
hand_rot = self._robot.data.body_link_quat_w[env_ids, self.hand_link_idx]
drawer_pos = self._cabinet.data.body_link_pos_w[env_ids, self.drawer_link_idx]
drawer_rot = self._cabinet.data.body_link_quat_w[env_ids, self.drawer_link_idx]
(
self.robot_grasp_rot[env_ids],
self.robot_grasp_pos[env_ids],
......
......@@ -221,7 +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_state_to_sim(object_default_state, env_ids)
self.object.write_root_link_pose_to_sim(object_default_state[:, :7], env_ids)
self.object.write_root_com_velocity_to_sim(object_default_state[:, 7:], env_ids)
# reset hand
delta_max = self.hand_dof_upper_limits[env_ids] - self.hand.data.default_joint_pos[env_ids]
......@@ -260,22 +261,22 @@ class InHandManipulationEnv(DirectRLEnv):
def _compute_intermediate_values(self):
# data for hand
self.fingertip_pos = self.hand.data.body_pos_w[:, self.finger_bodies]
self.fingertip_rot = self.hand.data.body_quat_w[:, self.finger_bodies]
self.fingertip_pos = self.hand.data.body_link_pos_w[:, self.finger_bodies]
self.fingertip_rot = self.hand.data.body_link_quat_w[:, self.finger_bodies]
self.fingertip_pos -= self.scene.env_origins.repeat((1, self.num_fingertips)).reshape(
self.num_envs, self.num_fingertips, 3
)
self.fingertip_velocities = self.hand.data.body_vel_w[:, self.finger_bodies]
self.fingertip_velocities = self.hand.data.body_com_vel_w[:, self.finger_bodies]
self.hand_dof_pos = self.hand.data.joint_pos
self.hand_dof_vel = self.hand.data.joint_vel
# data for object
self.object_pos = self.object.data.root_pos_w - self.scene.env_origins
self.object_rot = self.object.data.root_quat_w
self.object_velocities = self.object.data.root_vel_w
self.object_linvel = self.object.data.root_lin_vel_w
self.object_angvel = self.object.data.root_ang_vel_w
self.object_pos = self.object.data.root_link_pos_w - self.scene.env_origins
self.object_rot = self.object.data.root_link_quat_w
self.object_velocities = self.object.data.root_com_vel_w
self.object_linvel = self.object.data.root_com_lin_vel_w
self.object_angvel = self.object.data.root_com_ang_vel_w
def compute_reduced_observations(self):
# Per https://arxiv.org/pdf/1808.00177.pdf Table 2
......
......@@ -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_pos_w, self.robot.data.root_quat_w
self.velocity, self.ang_velocity = self.robot.data.root_lin_vel_w, self.robot.data.root_ang_vel_w
self.torso_position, self.torso_rotation = self.robot.data.root_link_pos_w, self.robot.data.root_link_quat_w
self.velocity, self.ang_velocity = self.robot.data.root_com_lin_vel_w, self.robot.data.root_com_ang_vel_w
self.dof_pos, self.dof_vel = self.robot.data.joint_pos, self.robot.data.joint_vel
(
......@@ -161,8 +161,8 @@ 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_pose_to_sim(default_root_state[:, :7], env_ids)
self.robot.write_root_velocity_to_sim(default_root_state[:, 7:], env_ids)
self.robot.write_root_link_pose_to_sim(default_root_state[:, :7], env_ids)
self.robot.write_root_com_velocity_to_sim(default_root_state[:, 7:], env_ids)
self.robot.write_joint_state_to_sim(joint_pos, joint_vel, None, env_ids)
to_target = self.targets[env_ids] - default_root_state[:, :3]
......
......@@ -154,12 +154,12 @@ class QuadcopterEnv(DirectRLEnv):
def _get_observations(self) -> dict:
desired_pos_b, _ = subtract_frame_transforms(
self._robot.data.root_state_w[:, :3], self._robot.data.root_state_w[:, 3:7], self._desired_pos_w
self._robot.data.root_link_state_w[:, :3], self._robot.data.root_link_state_w[:, 3:7], self._desired_pos_w
)
obs = torch.cat(
[
self._robot.data.root_lin_vel_b,
self._robot.data.root_ang_vel_b,
self._robot.data.root_com_lin_vel_b,
self._robot.data.root_com_ang_vel_b,
self._robot.data.projected_gravity_b,
desired_pos_b,
],
......@@ -169,9 +169,9 @@ class QuadcopterEnv(DirectRLEnv):
return observations
def _get_rewards(self) -> torch.Tensor:
lin_vel = torch.sum(torch.square(self._robot.data.root_lin_vel_b), dim=1)
ang_vel = torch.sum(torch.square(self._robot.data.root_ang_vel_b), dim=1)
distance_to_goal = torch.linalg.norm(self._desired_pos_w - self._robot.data.root_pos_w, dim=1)
lin_vel = torch.sum(torch.square(self._robot.data.root_com_lin_vel_b), dim=1)
ang_vel = torch.sum(torch.square(self._robot.data.root_com_ang_vel_b), dim=1)
distance_to_goal = torch.linalg.norm(self._desired_pos_w - self._robot.data.root_link_pos_w, dim=1)
distance_to_goal_mapped = 1 - torch.tanh(distance_to_goal / 0.8)
rewards = {
"lin_vel": lin_vel * self.cfg.lin_vel_reward_scale * self.step_dt,
......@@ -186,7 +186,9 @@ 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_pos_w[:, 2] < 0.1, self._robot.data.root_pos_w[:, 2] > 2.0)
died = torch.logical_or(
self._robot.data.root_link_pos_w[:, 2] < 0.1, self._robot.data.root_link_pos_w[:, 2] > 2.0
)
return died, time_out
def _reset_idx(self, env_ids: torch.Tensor | None):
......@@ -195,7 +197,7 @@ class QuadcopterEnv(DirectRLEnv):
# Logging
final_distance_to_goal = torch.linalg.norm(
self._desired_pos_w[env_ids] - self._robot.data.root_pos_w[env_ids], dim=1
self._desired_pos_w[env_ids] - self._robot.data.root_link_pos_w[env_ids], dim=1
).mean()
extras = dict()
for key in self._episode_sums.keys():
......@@ -226,8 +228,8 @@ 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_pose_to_sim(default_root_state[:, :7], env_ids)
self._robot.write_root_velocity_to_sim(default_root_state[:, 7:], env_ids)
self._robot.write_root_link_pose_to_sim(default_root_state[:, :7], env_ids)
self._robot.write_root_com_velocity_to_sim(default_root_state[:, 7:], env_ids)
self._robot.write_joint_state_to_sim(joint_pos, joint_vel, None, env_ids)
def _set_debug_vis_impl(self, debug_vis: bool):
......
......@@ -322,7 +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_state_to_sim(object_default_state, env_ids)
self.object.write_root_link_pose_to_sim(object_default_state[:, :7], env_ids)
self.object.write_root_com_velocity_to_sim(object_default_state[:, 7:], env_ids)
# reset right hand
delta_max = self.hand_dof_upper_limits[env_ids] - self.right_hand.data.default_joint_pos[env_ids]
......@@ -376,33 +377,33 @@ class ShadowHandOverEnv(DirectMARLEnv):
def _compute_intermediate_values(self):
# data for right hand
self.right_fingertip_pos = self.right_hand.data.body_pos_w[:, self.finger_bodies]
self.right_fingertip_rot = self.right_hand.data.body_quat_w[:, self.finger_bodies]
self.right_fingertip_pos = self.right_hand.data.body_link_pos_w[:, self.finger_bodies]
self.right_fingertip_rot = self.right_hand.data.body_link_quat_w[:, self.finger_bodies]
self.right_fingertip_pos -= self.scene.env_origins.repeat((1, self.num_fingertips)).reshape(
self.num_envs, self.num_fingertips, 3
)
self.right_fingertip_velocities = self.right_hand.data.body_vel_w[:, self.finger_bodies]
self.right_fingertip_velocities = self.right_hand.data.body_com_vel_w[:, self.finger_bodies]
self.right_hand_dof_pos = self.right_hand.data.joint_pos
self.right_hand_dof_vel = self.right_hand.data.joint_vel
# data for left hand
self.left_fingertip_pos = self.left_hand.data.body_pos_w[:, self.finger_bodies]
self.left_fingertip_rot = self.left_hand.data.body_quat_w[:, self.finger_bodies]
self.left_fingertip_pos = self.left_hand.data.body_link_pos_w[:, self.finger_bodies]
self.left_fingertip_rot = self.left_hand.data.body_link_quat_w[:, self.finger_bodies]
self.left_fingertip_pos -= self.scene.env_origins.repeat((1, self.num_fingertips)).reshape(
self.num_envs, self.num_fingertips, 3
)
self.left_fingertip_velocities = self.left_hand.data.body_vel_w[:, self.finger_bodies]
self.left_fingertip_velocities = self.left_hand.data.body_com_vel_w[:, self.finger_bodies]
self.left_hand_dof_pos = self.left_hand.data.joint_pos
self.left_hand_dof_vel = self.left_hand.data.joint_vel
# data for object
self.object_pos = self.object.data.root_pos_w - self.scene.env_origins
self.object_rot = self.object.data.root_quat_w
self.object_velocities = self.object.data.root_vel_w
self.object_linvel = self.object.data.root_lin_vel_w
self.object_angvel = self.object.data.root_ang_vel_w
self.object_pos = self.object.data.root_link_pos_w - self.scene.env_origins
self.object_rot = self.object.data.root_link_quat_w
self.object_velocities = self.object.data.root_com_vel_w
self.object_linvel = self.object.data.root_com_lin_vel_w
self.object_angvel = self.object.data.root_com_ang_vel_w
@torch.jit.script
......
......@@ -21,7 +21,7 @@ def base_yaw_roll(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityC
# extract the used quantities (to enable type-hinting)
asset: Articulation = env.scene[asset_cfg.name]
# extract euler angles (in world frame)
roll, _, yaw = math_utils.euler_xyz_from_quat(asset.data.root_quat_w)
roll, _, yaw = math_utils.euler_xyz_from_quat(asset.data.root_link_quat_w)
# normalize angle to [-pi, pi]
roll = torch.atan2(torch.sin(roll), torch.cos(roll))
yaw = torch.atan2(torch.sin(yaw), torch.cos(yaw))
......@@ -46,11 +46,11 @@ def base_heading_proj(
# extract the used quantities (to enable type-hinting)
asset: Articulation = env.scene[asset_cfg.name]
# compute desired heading direction
to_target_pos = torch.tensor(target_pos, device=env.device) - asset.data.root_pos_w[:, :3]
to_target_pos = torch.tensor(target_pos, device=env.device) - asset.data.root_link_pos_w[:, :3]
to_target_pos[:, 2] = 0.0
to_target_dir = math_utils.normalize(to_target_pos)
# compute base forward vector
heading_vec = math_utils.quat_rotate(asset.data.root_quat_w, asset.data.FORWARD_VEC_B)
heading_vec = math_utils.quat_rotate(asset.data.root_link_quat_w, asset.data.FORWARD_VEC_B)
# compute dot product between heading and target direction
heading_proj = torch.bmm(heading_vec.view(env.num_envs, 1, 3), to_target_dir.view(env.num_envs, 3, 1))
......@@ -64,10 +64,10 @@ def base_angle_to_target(
# extract the used quantities (to enable type-hinting)
asset: Articulation = env.scene[asset_cfg.name]
# compute desired heading direction
to_target_pos = torch.tensor(target_pos, device=env.device) - asset.data.root_pos_w[:, :3]
to_target_pos = torch.tensor(target_pos, device=env.device) - asset.data.root_link_pos_w[:, :3]
walk_target_angle = torch.atan2(to_target_pos[:, 1], to_target_pos[:, 0])
# compute base forward vector
_, _, yaw = math_utils.euler_xyz_from_quat(asset.data.root_quat_w)
_, _, yaw = math_utils.euler_xyz_from_quat(asset.data.root_link_quat_w)
# normalize angle to target to [-pi, pi]
angle_to_target = walk_target_angle - yaw
angle_to_target = torch.atan2(torch.sin(angle_to_target), torch.cos(angle_to_target))
......
......@@ -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_pos_w[env_ids, :3]
to_target_pos = target_pos - asset.data.root_link_pos_w[env_ids, :3]
# reward terms
self.potentials[env_ids] = -torch.norm(to_target_pos, p=2, dim=-1) / self._env.step_dt
self.prev_potentials[env_ids] = self.potentials[env_ids]
......@@ -68,7 +68,7 @@ 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_pos_w[:, :3]
to_target_pos = target_pos - asset.data.root_link_pos_w[:, :3]
to_target_pos[:, 2] = 0.0
# update history buffer and compute new potential
self.prev_potentials[:] = self.potentials[:]
......
......@@ -49,7 +49,7 @@ def air_time_reward(
t_min = torch.clip(t_max, max=mode_time)
stance_cmd_reward = torch.clip(current_contact_time - current_air_time, -mode_time, mode_time)
cmd = torch.norm(env.command_manager.get_command("base_velocity"), dim=1).unsqueeze(dim=1).expand(-1, 4)
body_vel = torch.linalg.norm(asset.data.root_lin_vel_b[:, :2], dim=1).unsqueeze(dim=1).expand(-1, 4)
body_vel = torch.linalg.norm(asset.data.root_com_lin_vel_b[:, :2], dim=1).unsqueeze(dim=1).expand(-1, 4)
reward = torch.where(
torch.logical_or(cmd > 0.0, body_vel > velocity_threshold),
torch.where(t_max < mode_time, t_min, 0),
......@@ -64,7 +64,7 @@ def base_angular_velocity_reward(env: ManagerBasedRLEnv, asset_cfg: SceneEntityC
asset: RigidObject = env.scene[asset_cfg.name]
# compute the error
target = env.command_manager.get_command("base_velocity")[:, 2]
ang_vel_error = torch.linalg.norm((target - asset.data.root_ang_vel_b[:, 2]).unsqueeze(1), dim=1)
ang_vel_error = torch.linalg.norm((target - asset.data.root_com_ang_vel_b[:, 2]).unsqueeze(1), dim=1)
return torch.exp(-ang_vel_error / std)
......@@ -76,7 +76,7 @@ def base_linear_velocity_reward(
asset: RigidObject = env.scene[asset_cfg.name]
# compute the error
target = env.command_manager.get_command("base_velocity")[:, :2]
lin_vel_error = torch.linalg.norm((target - asset.data.root_lin_vel_b[:, :2]), dim=1)
lin_vel_error = torch.linalg.norm((target - asset.data.root_com_lin_vel_b[:, :2]), dim=1)
# fixed 1.0 multiple for tracking below the ramp_at_vel value, then scale by the rate above
vel_cmd_magnitude = torch.linalg.norm(target, dim=1)
velocity_scaling_multiple = torch.clamp(1.0 + ramp_rate * (vel_cmd_magnitude - ramp_at_vel), min=1.0)
......@@ -148,7 +148,7 @@ 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_lin_vel_b[:, :2], dim=1)
body_vel = torch.linalg.norm(self.asset.data.root_com_lin_vel_b[:, :2], dim=1)
return torch.where(
torch.logical_or(cmd > 0.0, body_vel > self.velocity_threshold), sync_reward * async_reward, 0.0
)
......@@ -182,8 +182,10 @@ def foot_clearance_reward(
) -> torch.Tensor:
"""Reward the swinging feet for clearing a specified height off the ground"""
asset: RigidObject = env.scene[asset_cfg.name]
foot_z_target_error = torch.square(asset.data.body_pos_w[:, asset_cfg.body_ids, 2] - target_height)
foot_velocity_tanh = torch.tanh(tanh_mult * torch.norm(asset.data.body_lin_vel_w[:, asset_cfg.body_ids, :2], dim=2))
foot_z_target_error = torch.square(asset.data.body_link_pos_w[:, asset_cfg.body_ids, 2] - target_height)
foot_velocity_tanh = torch.tanh(
tanh_mult * torch.norm(asset.data.body_com_lin_vel_w[:, asset_cfg.body_ids, :2], dim=2)
)
reward = foot_z_target_error * foot_velocity_tanh
return torch.exp(-torch.sum(reward, dim=1) / std)
......@@ -217,8 +219,8 @@ def base_motion_penalty(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg) -> to
"""Penalize base vertical and roll/pitch velocity"""
# extract the used quantities (to enable type-hinting)
asset: RigidObject = env.scene[asset_cfg.name]
return 0.8 * torch.square(asset.data.root_lin_vel_b[:, 2]) + 0.2 * torch.sum(
torch.abs(asset.data.root_ang_vel_b[:, :2]), dim=1
return 0.8 * torch.square(asset.data.root_com_lin_vel_b[:, 2]) + 0.2 * torch.sum(
torch.abs(asset.data.root_com_ang_vel_b[:, :2]), dim=1
)
......@@ -243,7 +245,7 @@ def foot_slip_penalty(
# check if contact force is above threshold
net_contact_forces = contact_sensor.data.net_forces_w_history
is_contact = torch.max(torch.norm(net_contact_forces[:, :, sensor_cfg.body_ids], dim=-1), dim=1)[0] > threshold
foot_planar_velocity = torch.linalg.norm(asset.data.body_lin_vel_w[:, asset_cfg.body_ids, :2], dim=2)
foot_planar_velocity = torch.linalg.norm(asset.data.body_com_lin_vel_w[:, asset_cfg.body_ids, :2], dim=2)
reward = is_contact * foot_planar_velocity
return torch.sum(reward, dim=1)
......@@ -263,7 +265,7 @@ def joint_position_penalty(
# extract the used quantities (to enable type-hinting)
asset: Articulation = env.scene[asset_cfg.name]
cmd = torch.linalg.norm(env.command_manager.get_command("base_velocity"), dim=1)
body_vel = torch.linalg.norm(asset.data.root_lin_vel_b[:, :2], dim=1)
body_vel = torch.linalg.norm(asset.data.root_com_lin_vel_b[:, :2], dim=1)
reward = torch.linalg.norm((asset.data.joint_pos - asset.data.default_joint_pos), dim=1)
return torch.where(torch.logical_or(cmd > 0.0, body_vel > velocity_threshold), reward, stand_still_scale * reward)
......
......@@ -43,7 +43,7 @@ def terrain_levels_vel(
terrain: TerrainImporter = env.scene.terrain
command = env.command_manager.get_command("base_velocity")
# compute the distance the robot walked
distance = torch.norm(asset.data.root_pos_w[env_ids, :2] - env.scene.env_origins[env_ids, :2], dim=1)
distance = torch.norm(asset.data.root_link_pos_w[env_ids, :2] - env.scene.env_origins[env_ids, :2], dim=1)
# robots that walked far enough progress to harder terrains
move_up = distance > terrain.cfg.terrain_generator.size[0] / 2
# robots that walked less than half of their required distance go to simpler terrains
......
......@@ -77,7 +77,8 @@ def feet_slide(env, sensor_cfg: SceneEntityCfg, asset_cfg: SceneEntityCfg = Scen
contact_sensor: ContactSensor = env.scene.sensors[sensor_cfg.name]
contacts = contact_sensor.data.net_forces_w_history[:, :, sensor_cfg.body_ids, :].norm(dim=-1).max(dim=1)[0] > 1.0
asset = env.scene[asset_cfg.name]
body_vel = asset.data.body_lin_vel_w[:, sensor_cfg.body_ids, :2]
body_vel = asset.data.body_com_lin_vel_w[:, asset_cfg.body_ids, :2]
reward = torch.sum(body_vel.norm(dim=-1) * contacts, dim=1)
return reward
......@@ -88,7 +89,7 @@ def track_lin_vel_xy_yaw_frame_exp(
"""Reward tracking of linear velocity commands (xy axes) in the gravity aligned robot frame using exponential kernel."""
# extract the used quantities (to enable type-hinting)
asset = env.scene[asset_cfg.name]
vel_yaw = quat_rotate_inverse(yaw_quat(asset.data.root_quat_w), asset.data.root_lin_vel_w[:, :3])
vel_yaw = quat_rotate_inverse(yaw_quat(asset.data.root_link_quat_w), asset.data.root_com_lin_vel_w[:, :3])
lin_vel_error = torch.sum(
torch.square(env.command_manager.get_command(command_name)[:, :2] - vel_yaw[:, :2]), dim=1
)
......@@ -101,5 +102,7 @@ def track_ang_vel_z_world_exp(
"""Reward tracking of angular velocity commands (yaw) in world frame using exponential kernel."""
# extract the used quantities (to enable type-hinting)
asset = env.scene[asset_cfg.name]
ang_vel_error = torch.square(env.command_manager.get_command(command_name)[:, 2] - asset.data.root_ang_vel_w[:, 2])
ang_vel_error = torch.square(
env.command_manager.get_command(command_name)[:, 2] - asset.data.root_com_ang_vel_w[:, 2]
)
return torch.exp(-ang_vel_error / std**2)
......@@ -45,8 +45,8 @@ def terrain_out_of_bounds(
asset: RigidObject = env.scene[asset_cfg.name]
# check if the agent is out of bounds
x_out_of_bounds = torch.abs(asset.data.root_pos_w[:, 0]) > 0.5 * map_width - distance_buffer
y_out_of_bounds = torch.abs(asset.data.root_pos_w[:, 1]) > 0.5 * map_height - distance_buffer
x_out_of_bounds = torch.abs(asset.data.root_link_pos_w[:, 0]) > 0.5 * map_width - distance_buffer
y_out_of_bounds = torch.abs(asset.data.root_link_pos_w[:, 1]) > 0.5 * map_height - distance_buffer
return torch.logical_or(x_out_of_bounds, y_out_of_bounds)
else:
raise ValueError("Received unsupported terrain type, must be either 'plane' or 'generator'.")
......@@ -21,7 +21,7 @@ def rel_ee_object_distance(env: ManagerBasedRLEnv) -> torch.Tensor:
ee_tf_data: FrameTransformerData = env.scene["ee_frame"].data
object_data: ArticulationData = env.scene["object"].data
return object_data.root_pos_w - ee_tf_data.target_pos_w[..., 0, :]
return object_data.root_link_pos_w - ee_tf_data.target_pos_w[..., 0, :]
def rel_ee_drawer_distance(env: ManagerBasedRLEnv) -> torch.Tensor:
......
......@@ -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_quat_w, self.quat_command_w
self.object.data.root_link_quat_w, self.quat_command_w
)
# -- compute the position error
self.metrics["position_error"] = torch.norm(self.object.data.root_pos_w - self.pos_command_w, dim=1)
self.metrics["position_error"] = torch.norm(self.object.data.root_link_pos_w - self.pos_command_w, dim=1)
# -- compute the number of consecutive successes
successes = self.metrics["orientation_error"] < self.cfg.orientation_success_threshold
self.metrics["consecutive_success"] += successes.float()
......
......@@ -30,7 +30,7 @@ def goal_quat_diff(
# obtain the orientations
goal_quat_w = command_term.command[:, 3:7]
asset_quat_w = asset.data.root_quat_w
asset_quat_w = asset.data.root_link_quat_w
# compute quaternion difference
quat = math_utils.quat_mul(asset_quat_w, math_utils.quat_conjugate(goal_quat_w))
......
......@@ -39,7 +39,7 @@ def success_bonus(
# obtain the threshold for the orientation error
threshold = command_term.cfg.orientation_success_threshold
# calculate the orientation error
dtheta = math_utils.quat_error_magnitude(asset.data.root_quat_w, goal_quat_w)
dtheta = math_utils.quat_error_magnitude(asset.data.root_link_quat_w, goal_quat_w)
return dtheta <= threshold
......@@ -63,7 +63,7 @@ def track_pos_l2(
# obtain the goal position
goal_pos_e = command_term.command[:, 0:3]
# obtain the object position in the environment frame
object_pos_e = asset.data.root_pos_w - env.scene.env_origins
object_pos_e = asset.data.root_link_pos_w - env.scene.env_origins
return torch.norm(goal_pos_e - object_pos_e, p=2, dim=-1)
......@@ -91,6 +91,6 @@ def track_orientation_inv_l2(
# obtain the goal orientation
goal_quat_w = command_term.command[:, 3:7]
# calculate the orientation error
dtheta = math_utils.quat_error_magnitude(asset.data.root_quat_w, goal_quat_w)
dtheta = math_utils.quat_error_magnitude(asset.data.root_link_quat_w, goal_quat_w)
return 1.0 / (dtheta + rot_eps)
......@@ -50,7 +50,7 @@ def object_away_from_goal(
asset = env.scene[object_cfg.name]
# object pos
asset_pos_e = asset.data.root_pos_w - env.scene.env_origins
asset_pos_e = asset.data.root_link_pos_w - env.scene.env_origins
goal_pos_e = command_term.command[:, :3]
return torch.norm(asset_pos_e - goal_pos_e, p=2, dim=1) > threshold
......@@ -78,6 +78,6 @@ def object_away_from_robot(
object = env.scene[object_cfg.name]
# compute distance
dist = torch.norm(robot.data.root_pos_w - object.data.root_pos_w, dim=1)
dist = torch.norm(robot.data.root_link_pos_w - object.data.root_link_pos_w, dim=1)
return dist > threshold
......@@ -24,8 +24,8 @@ def object_position_in_robot_root_frame(
"""The position of the object in the robot's root frame."""
robot: RigidObject = env.scene[robot_cfg.name]
object: RigidObject = env.scene[object_cfg.name]
object_pos_w = object.data.root_pos_w[:, :3]
object_pos_w = object.data.root_link_pos_w[:, :3]
object_pos_b, _ = subtract_frame_transforms(
robot.data.root_state_w[:, :3], robot.data.root_state_w[:, 3:7], object_pos_w
robot.data.root_link_state_w[:, :3], robot.data.root_link_state_w[:, 3:7], object_pos_w
)
return object_pos_b
......@@ -22,7 +22,7 @@ def object_is_lifted(
) -> torch.Tensor:
"""Reward the agent for lifting the object above the minimal height."""
object: RigidObject = env.scene[object_cfg.name]
return torch.where(object.data.root_pos_w[:, 2] > minimal_height, 1.0, 0.0)
return torch.where(object.data.root_link_pos_w[:, 2] > minimal_height, 1.0, 0.0)
def object_ee_distance(
......@@ -36,7 +36,7 @@ def object_ee_distance(
object: RigidObject = env.scene[object_cfg.name]
ee_frame: FrameTransformer = env.scene[ee_frame_cfg.name]
# Target object position: (num_envs, 3)
cube_pos_w = object.data.root_pos_w
cube_pos_w = object.data.root_link_pos_w
# End-effector position: (num_envs, 3)
ee_w = ee_frame.data.target_pos_w[..., 0, :]
# Distance of the end-effector to the object: (num_envs,)
......@@ -60,8 +60,10 @@ def object_goal_distance(
command = env.command_manager.get_command(command_name)
# compute the desired position in the world frame
des_pos_b = command[:, :3]
des_pos_w, _ = combine_frame_transforms(robot.data.root_state_w[:, :3], robot.data.root_state_w[:, 3:7], des_pos_b)
des_pos_w, _ = combine_frame_transforms(
robot.data.root_link_state_w[:, :3], robot.data.root_link_state_w[:, 3:7], des_pos_b
)
# distance of the end-effector to the object: (num_envs,)
distance = torch.norm(des_pos_w - object.data.root_pos_w[:, :3], dim=1)
distance = torch.norm(des_pos_w - object.data.root_link_pos_w[:, :3], dim=1)
# rewarded if the object is lifted above the threshold
return (object.data.root_pos_w[:, 2] > minimal_height) * (1 - torch.tanh(distance / std))
return (object.data.root_link_pos_w[:, 2] > minimal_height) * (1 - torch.tanh(distance / std))
......@@ -45,9 +45,11 @@ def object_reached_goal(
command = env.command_manager.get_command(command_name)
# compute the desired position in the world frame
des_pos_b = command[:, :3]
des_pos_w, _ = combine_frame_transforms(robot.data.root_state_w[:, :3], robot.data.root_state_w[:, 3:7], des_pos_b)
des_pos_w, _ = combine_frame_transforms(
robot.data.root_link_state_w[:, :3], robot.data.root_link_state_w[:, 3:7], des_pos_b
)
# distance of the end-effector to the object: (num_envs,)
distance = torch.norm(des_pos_w - object.data.root_pos_w[:, :3], dim=1)
distance = torch.norm(des_pos_w - object.data.root_link_pos_w[:, :3], dim=1)
# rewarded if the object is lifted above the threshold
return distance < threshold
......@@ -28,8 +28,10 @@ def position_command_error(env: ManagerBasedRLEnv, command_name: str, asset_cfg:
command = env.command_manager.get_command(command_name)
# obtain the desired and current positions
des_pos_b = command[:, :3]
des_pos_w, _ = combine_frame_transforms(asset.data.root_state_w[:, :3], asset.data.root_state_w[:, 3:7], des_pos_b)
curr_pos_w = asset.data.body_state_w[:, asset_cfg.body_ids[0], :3] # type: ignore
des_pos_w, _ = combine_frame_transforms(
asset.data.root_link_state_w[:, :3], asset.data.root_link_state_w[:, 3:7], des_pos_b
)
curr_pos_w = asset.data.body_link_state_w[:, asset_cfg.body_ids[0], :3] # type: ignore
return torch.norm(curr_pos_w - des_pos_w, dim=1)
......@@ -46,8 +48,10 @@ def position_command_error_tanh(
command = env.command_manager.get_command(command_name)
# obtain the desired and current positions
des_pos_b = command[:, :3]
des_pos_w, _ = combine_frame_transforms(asset.data.root_state_w[:, :3], asset.data.root_state_w[:, 3:7], des_pos_b)
curr_pos_w = asset.data.body_state_w[:, asset_cfg.body_ids[0], :3] # type: ignore
des_pos_w, _ = combine_frame_transforms(
asset.data.root_link_state_w[:, :3], asset.data.root_link_state_w[:, 3:7], des_pos_b
)
curr_pos_w = asset.data.body_link_state_w[:, asset_cfg.body_ids[0], :3] # type: ignore
distance = torch.norm(curr_pos_w - des_pos_w, dim=1)
return 1 - torch.tanh(distance / std)
......@@ -64,6 +68,6 @@ def orientation_command_error(env: ManagerBasedRLEnv, command_name: str, asset_c
command = env.command_manager.get_command(command_name)
# obtain the desired and current orientations
des_quat_b = command[:, 3:7]
des_quat_w = quat_mul(asset.data.root_state_w[:, 3:7], des_quat_b)
curr_quat_w = asset.data.body_state_w[:, asset_cfg.body_ids[0], 3:7] # type: ignore
des_quat_w = quat_mul(asset.data.root_link_state_w[:, 3:7], des_quat_b)
curr_quat_w = asset.data.body_link_state_w[:, asset_cfg.body_ids[0], 3:7] # type: ignore
return quat_error_magnitude(curr_quat_w, des_quat_w)
......@@ -130,10 +130,10 @@ def randomize_object_pose(
pose_tensor = torch.tensor([pose_list[i]], device=env.device)
positions = pose_tensor[:, 0:3] + env.scene.env_origins[cur_env, 0:3]
orientations = math_utils.quat_from_euler_xyz(pose_tensor[:, 3], pose_tensor[:, 4], pose_tensor[:, 5])
asset.write_root_pose_to_sim(
asset.write_root_link_pose_to_sim(
torch.cat([positions, orientations], dim=-1), env_ids=torch.tensor([cur_env], device=env.device)
)
asset.write_root_velocity_to_sim(
asset.write_root_com_velocity_to_sim(
torch.zeros(1, 6, device=env.device), env_ids=torch.tensor([cur_env], device=env.device)
)
......
......@@ -27,7 +27,7 @@ def cube_positions_in_world_frame(
cube_2: RigidObject = env.scene[cube_2_cfg.name]
cube_3: RigidObject = env.scene[cube_3_cfg.name]
return torch.cat((cube_1.data.root_pos_w, cube_2.data.root_pos_w, cube_3.data.root_pos_w), dim=1)
return torch.cat((cube_1.data.root_link_pos_w, cube_2.data.root_link_pos_w, cube_3.data.root_link_pos_w), dim=1)
def instance_randomize_cube_positions_in_world_frame(
......@@ -69,7 +69,7 @@ def cube_orientations_in_world_frame(
cube_2: RigidObject = env.scene[cube_2_cfg.name]
cube_3: RigidObject = env.scene[cube_3_cfg.name]
return torch.cat((cube_1.data.root_quat_w, cube_2.data.root_quat_w, cube_3.data.root_quat_w), dim=1)
return torch.cat((cube_1.data.root_link_quat_w, cube_2.data.root_link_quat_w, cube_3.data.root_link_quat_w), dim=1)
def instance_randomize_cube_orientations_in_world_frame(
......@@ -127,14 +127,14 @@ def object_obs(
cube_3: RigidObject = env.scene[cube_3_cfg.name]
ee_frame: FrameTransformer = env.scene[ee_frame_cfg.name]
cube_1_pos_w = cube_1.data.root_pos_w
cube_1_quat_w = cube_1.data.root_quat_w
cube_1_pos_w = cube_1.data.root_link_pos_w
cube_1_quat_w = cube_1.data.root_link_quat_w
cube_2_pos_w = cube_2.data.root_pos_w
cube_2_quat_w = cube_2.data.root_quat_w
cube_2_pos_w = cube_2.data.root_link_pos_w
cube_2_quat_w = cube_2.data.root_link_quat_w
cube_3_pos_w = cube_3.data.root_pos_w
cube_3_quat_w = cube_3.data.root_quat_w
cube_3_pos_w = cube_3.data.root_link_pos_w
cube_3_quat_w = cube_3.data.root_link_quat_w
ee_pos_w = ee_frame.data.target_pos_w[:, 0, :]
gripper_to_cube_1 = cube_1_pos_w - ee_pos_w
......@@ -279,7 +279,7 @@ def object_grasped(
ee_frame: FrameTransformer = env.scene[ee_frame_cfg.name]
object: RigidObject = env.scene[object_cfg.name]
object_pos = object.data.root_pos_w
object_pos = object.data.root_link_pos_w
end_effector_pos = ee_frame.data.target_pos_w[:, 0, :]
pose_diff = torch.linalg.vector_norm(object_pos - end_effector_pos, dim=1)
......@@ -310,7 +310,7 @@ def object_stacked(
upper_object: RigidObject = env.scene[upper_object_cfg.name]
lower_object: RigidObject = env.scene[lower_object_cfg.name]
pos_diff = upper_object.data.root_pos_w - lower_object.data.root_pos_w
pos_diff = upper_object.data.root_link_pos_w - lower_object.data.root_link_pos_w
height_dist = torch.linalg.vector_norm(pos_diff[:, 2:], dim=1)
xy_dist = torch.linalg.vector_norm(pos_diff[:, :2], dim=1)
......
......@@ -39,8 +39,8 @@ def cubes_stacked(
cube_2: RigidObject = env.scene[cube_2_cfg.name]
cube_3: RigidObject = env.scene[cube_3_cfg.name]
pos_diff_c12 = cube_1.data.root_pos_w - cube_2.data.root_pos_w
pos_diff_c23 = cube_2.data.root_pos_w - cube_3.data.root_pos_w
pos_diff_c12 = cube_1.data.root_link_pos_w - cube_2.data.root_link_pos_w
pos_diff_c23 = cube_2.data.root_link_pos_w - cube_3.data.root_link_pos_w
# Compute cube position difference in x-y plane
xy_dist_c12 = torch.norm(pos_diff_c12[:, :2], dim=1)
......
......@@ -128,11 +128,11 @@ class PreTrainedPolicyAction(ActionTerm):
return
# get marker location
# -- base state
base_pos_w = self.robot.data.root_pos_w.clone()
base_pos_w = self.robot.data.root_link_pos_w.clone()
base_pos_w[:, 2] += 0.5
# -- resolve the scales and quaternions
vel_des_arrow_scale, vel_des_arrow_quat = self._resolve_xy_velocity_to_arrow(self.raw_actions[:, :2])
vel_arrow_scale, vel_arrow_quat = self._resolve_xy_velocity_to_arrow(self.robot.data.root_lin_vel_b[:, :2])
vel_arrow_scale, vel_arrow_quat = self._resolve_xy_velocity_to_arrow(self.robot.data.root_com_lin_vel_b[:, :2])
# display markers
self.base_vel_goal_visualizer.visualize(base_pos_w, vel_des_arrow_quat, vel_des_arrow_scale)
self.base_vel_visualizer.visualize(base_pos_w, vel_arrow_quat, vel_arrow_scale)
......@@ -153,7 +153,7 @@ 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_quat_w
base_quat_w = self.robot.data.root_link_quat_w
arrow_quat = math_utils.quat_mul(base_quat_w, arrow_quat)
return arrow_scale, arrow_quat
......
......@@ -115,7 +115,8 @@ def run_simulator(sim: sim_utils.SimulationContext, scene: InteractiveScene):
# if this is not done, then the robots will be spawned at the (0, 0, 0) of the simulation world
root_state = robot.data.default_root_state.clone()
root_state[:, :3] += scene.env_origins
robot.write_root_state_to_sim(root_state)
robot.write_root_link_pose_to_sim(root_state[:, :7])
robot.write_root_com_velocity_to_sim(root_state[:, 7:])
# set joint positions with some noise
joint_pos, joint_vel = robot.data.default_joint_pos.clone(), robot.data.default_joint_vel.clone()
joint_pos += torch.rand_like(joint_pos) * 0.1
......
......@@ -179,7 +179,8 @@ def run_simulator(sim: sim_utils.SimulationContext, entities: dict[str, Articula
# root state
root_state = robot.data.default_root_state.clone()
root_state[:, :3] += origins[index]
robot.write_root_state_to_sim(root_state)
robot.write_root_link_pose_to_sim(root_state[:, :7])
robot.write_root_com_velocity_to_sim(root_state[:, 7:])
# set joint positions
joint_pos, joint_vel = robot.data.default_joint_pos.clone(), robot.data.default_joint_vel.clone()
robot.write_joint_state_to_sim(joint_pos, joint_vel)
......
......@@ -97,7 +97,8 @@ def main():
robot.write_joint_state_to_sim(joint_pos, joint_vel)
root_state = robot.data.default_root_state.clone()
root_state[:, :3] += origins[index]
robot.write_root_state_to_sim(root_state)
robot.write_root_link_pose_to_sim(root_state[:, :7])
robot.write_root_com_velocity_to_sim(root_state[:, 7:])
robot.reset()
# reset command
print(">>>>>>>> Reset!")
......
......@@ -187,7 +187,8 @@ def run_simulator(sim: sim_utils.SimulationContext, scene: InteractiveScene):
# if this is not done, then the robots will be spawned at the (0, 0, 0) of the simulation world
root_state = scene["robot"].data.default_root_state.clone()
root_state[:, :3] += scene.env_origins
scene["robot"].write_root_state_to_sim(root_state)
scene["robot"].write_root_link_pose_to_sim(root_state[:, :7])
scene["robot"].write_root_com_velocity_to_sim(root_state[:, 7:])
# set joint positions with some noise
joint_pos, joint_vel = (
scene["robot"].data.default_joint_pos.clone(),
......
......@@ -113,7 +113,8 @@ def run_simulator(sim: sim_utils.SimulationContext, entities: dict[str, Articula
# root state
root_state = robot.data.default_root_state.clone()
root_state[:, :3] += origins[index]
robot.write_root_state_to_sim(root_state)
robot.write_root_link_pose_to_sim(root_state[:, :7])
robot.write_root_com_velocity_to_sim(root_state[:, 7:])
# joint state
joint_pos, joint_vel = robot.data.default_joint_pos.clone(), robot.data.default_joint_vel.clone()
robot.write_joint_state_to_sim(joint_pos, joint_vel)
......
......@@ -240,16 +240,19 @@ def run_simulator(sim: SimulationContext, scene: InteractiveScene):
# object
root_state = rigid_object.data.default_root_state.clone()
root_state[:, :3] += scene.env_origins
rigid_object.write_root_state_to_sim(root_state)
rigid_object.write_root_link_pose_to_sim(root_state[:, :7])
rigid_object.write_root_com_velocity_to_sim(root_state[:, 7:])
# object collection
object_state = rigid_object_collection.data.default_object_state.clone()
object_state[..., :3] += scene.env_origins.unsqueeze(1)
rigid_object_collection.write_object_state_to_sim(object_state)
rigid_object_collection.write_object_link_pose_to_sim(object_state[..., :7])
rigid_object_collection.write_object_com_velocity_to_sim(object_state[..., 7:])
# robot
# -- root state
root_state = robot.data.default_root_state.clone()
root_state[:, :3] += scene.env_origins
robot.write_root_state_to_sim(root_state)
robot.write_root_link_pose_to_sim(root_state[:, :7])
robot.write_root_com_velocity_to_sim(root_state[:, 7:])
# -- joint state
joint_pos, joint_vel = robot.data.default_joint_pos.clone(), robot.data.default_joint_vel.clone()
robot.write_joint_state_to_sim(joint_pos, joint_vel)
......
......@@ -91,8 +91,8 @@ def main():
# reset dof state
joint_pos, joint_vel = robot.data.default_joint_pos, robot.data.default_joint_vel
robot.write_joint_state_to_sim(joint_pos, joint_vel)
robot.write_root_pose_to_sim(robot.data.default_root_state[:, :7])
robot.write_root_velocity_to_sim(robot.data.default_root_state[:, 7:])
robot.write_root_link_pose_to_sim(robot.data.default_root_state[:, :7])
robot.write_root_com_velocity_to_sim(robot.data.default_root_state[:, 7:])
robot.reset()
# reset command
print(">>>>>>>> Reset!")
......
......@@ -142,7 +142,8 @@ def run_simulator(sim: sim_utils.SimulationContext, entities: dict[str, Articula
# root state
root_state = robot.data.default_root_state.clone()
root_state[:, :3] += origins[index]
robot.write_root_state_to_sim(root_state)
robot.write_root_link_pose_to_sim(root_state[:, :7])
robot.write_root_com_velocity_to_sim(root_state[:, 7:])
# joint state
joint_pos, joint_vel = robot.data.default_joint_pos.clone(), robot.data.default_joint_vel.clone()
robot.write_joint_state_to_sim(joint_pos, joint_vel)
......
......@@ -109,7 +109,8 @@ def run_simulator(sim: sim_utils.SimulationContext, scene: InteractiveScene):
# if this is not done, then the robots will be spawned at the (0, 0, 0) of the simulation world
root_state = scene["robot"].data.default_root_state.clone()
root_state[:, :3] += scene.env_origins
scene["robot"].write_root_state_to_sim(root_state)
scene["robot"].write_root_link_pose_to_sim(root_state[:, :7])
scene["robot"].write_root_com_velocity_to_sim(root_state[:, 7:])
# set joint positions with some noise
joint_pos, joint_vel = (
scene["robot"].data.default_joint_pos.clone(),
......
......@@ -105,7 +105,8 @@ def run_simulator(sim: sim_utils.SimulationContext, scene: InteractiveScene):
# if this is not done, then the robots will be spawned at the (0, 0, 0) of the simulation world
root_state = scene["robot"].data.default_root_state.clone()
root_state[:, :3] += scene.env_origins
scene["robot"].write_root_state_to_sim(root_state)
scene["robot"].write_root_link_pose_to_sim(root_state[:, :7])
scene["robot"].write_root_com_velocity_to_sim(root_state[:, 7:])
# set joint positions with some noise
joint_pos, joint_vel = (
scene["robot"].data.default_joint_pos.clone(),
......
......@@ -63,6 +63,7 @@ class RaycasterSensorSceneCfg(InteractiveSceneCfg):
update_period=1 / 60,
offset=RayCasterCfg.OffsetCfg(pos=(0, 0, 0.5)),
mesh_prim_paths=["/World/Ground"],
attach_yaw_only=True,
pattern_cfg=patterns.LidarPatternCfg(
channels=100, vertical_fov_range=[-90, 90], horizontal_fov_range=[-90, 90], horizontal_res=1.0
),
......@@ -92,7 +93,8 @@ def run_simulator(sim: sim_utils.SimulationContext, scene: InteractiveScene):
# if this is not done, then the robots will be spawned at the (0, 0, 0) of the simulation world
root_state = scene["robot"].data.default_root_state.clone()
root_state[:, :3] += scene.env_origins
scene["robot"].write_root_state_to_sim(root_state)
scene["robot"].write_root_link_pose_to_sim(root_state[:, :7])
scene["robot"].write_root_com_velocity_to_sim(root_state[:, 7:])
# set joint positions with some noise
joint_pos, joint_vel = (
scene["robot"].data.default_joint_pos.clone(),
......
......@@ -293,7 +293,7 @@ def main():
tcp_rest_orientation = ee_frame_sensor.data.target_quat_w[..., 0, :].clone()
# -- object frame
object_data: RigidObjectData = env.unwrapped.scene["object"].data
object_position = object_data.root_pos_w - env.unwrapped.scene.env_origins
object_position = object_data.root_link_pos_w - env.unwrapped.scene.env_origins
# -- target object frame
desired_position = env.unwrapped.command_manager.get_command("object_pose")[..., :3]
......
......@@ -94,7 +94,8 @@ def run_simulator(sim: sim_utils.SimulationContext, entities: dict[str, Articula
# if this is not done, then the robots will be spawned at the (0, 0, 0) of the simulation world
root_state = robot.data.default_root_state.clone()
root_state[:, :3] += origins
robot.write_root_state_to_sim(root_state)
robot.write_root_link_pose_to_sim(root_state[:, :7])
robot.write_root_com_velocity_to_sim(root_state[:, 7:])
# set joint positions with some noise
joint_pos, joint_vel = robot.data.default_joint_pos.clone(), robot.data.default_joint_vel.clone()
joint_pos += torch.rand_like(joint_pos) * 0.1
......
......@@ -103,7 +103,8 @@ def run_simulator(sim: sim_utils.SimulationContext, entities: dict[str, RigidObj
radius=0.1, h_range=(0.25, 0.5), size=cone_object.num_instances, device=cone_object.device
)
# write root state to simulation
cone_object.write_root_state_to_sim(root_state)
cone_object.write_root_link_pose_to_sim(root_state[:, :7])
cone_object.write_root_com_velocity_to_sim(root_state[:, 7:])
# reset buffers
cone_object.reset()
print("----------------------------------------")
......@@ -119,7 +120,7 @@ def run_simulator(sim: sim_utils.SimulationContext, entities: dict[str, RigidObj
cone_object.update(sim_dt)
# print the root position
if count % 50 == 0:
print(f"Root position (in world): {cone_object.data.root_state_w[:, :3]}")
print(f"Root position (in world): {cone_object.data.root_link_state_w[:, :3]}")
def main():
......
......@@ -83,7 +83,8 @@ def run_simulator(sim: sim_utils.SimulationContext, scene: InteractiveScene):
# if this is not done, then the robots will be spawned at the (0, 0, 0) of the simulation world
root_state = robot.data.default_root_state.clone()
root_state[:, :3] += scene.env_origins
robot.write_root_state_to_sim(root_state)
robot.write_root_link_pose_to_sim(root_state[:, :7])
robot.write_root_com_velocity_to_sim(root_state[:, 7:])
# set joint positions with some noise
joint_pos, joint_vel = robot.data.default_joint_pos.clone(), robot.data.default_joint_vel.clone()
joint_pos += torch.rand_like(joint_pos) * 0.1
......
......@@ -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_pos_w - self._env.scene.env_origins)
vel_error = -self._asset.data.root_lin_vel_w
pos_error = self._processed_actions - (self._asset.data.root_link_pos_w - self._env.scene.env_origins)
vel_error = -self._asset.data.root_com_lin_vel_w
# set velocity targets
self._vel_command[:, :3] = self.p_gain * pos_error + self.d_gain * vel_error
self._asset.write_root_velocity_to_sim(self._vel_command)
self._asset.write_root_com_velocity_to_sim(self._vel_command)
@configclass
......@@ -149,7 +149,7 @@ def base_position(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg) -> torch.Tens
"""Root linear velocity in the asset's root frame."""
# extract the used quantities (to enable type-hinting)
asset: RigidObject = env.scene[asset_cfg.name]
return asset.data.root_pos_w - env.scene.env_origins
return asset.data.root_link_pos_w - env.scene.env_origins
##
......
......@@ -113,7 +113,8 @@ def run_simulator(sim: sim_utils.SimulationContext, scene: InteractiveScene):
# if this is not done, then the robots will be spawned at the (0, 0, 0) of the simulation world
root_state = scene["robot"].data.default_root_state.clone()
root_state[:, :3] += scene.env_origins
scene["robot"].write_root_state_to_sim(root_state)
scene["robot"].write_root_link_pose_to_sim(root_state[:, :7])
scene["robot"].write_root_com_velocity_to_sim(root_state[:, 7:])
# set joint positions with some noise
joint_pos, joint_vel = (
scene["robot"].data.default_joint_pos.clone(),
......
......@@ -109,7 +109,8 @@ def run_simulator(sim: sim_utils.SimulationContext, scene_entities: dict):
# Reset the scene
if step_count % 250 == 0:
# reset the balls
balls.write_root_state_to_sim(ball_default_state)
balls.write_root_link_pose_to_sim(ball_default_state[:, :7])
balls.write_root_com_velocity_to_sim(ball_default_state[:, 7:])
# reset the sensor
ray_caster.reset()
# reset the counter
......
......@@ -160,8 +160,8 @@ def run_simulator(sim: sim_utils.SimulationContext, scene: InteractiveScene):
else:
# obtain quantities from simulation
jacobian = robot.root_physx_view.get_jacobians()[:, ee_jacobi_idx, :, robot_entity_cfg.joint_ids]
ee_pose_w = robot.data.body_state_w[:, robot_entity_cfg.body_ids[0], 0:7]
root_pose_w = robot.data.root_state_w[:, 0:7]
ee_pose_w = robot.data.body_link_state_w[:, robot_entity_cfg.body_ids[0], 0:7]
root_pose_w = robot.data.root_link_state_w[:, 0:7]
joint_pos = robot.data.joint_pos[:, robot_entity_cfg.joint_ids]
# compute frame in root frame
ee_pos_b, ee_quat_b = subtract_frame_transforms(
......@@ -181,7 +181,7 @@ def run_simulator(sim: sim_utils.SimulationContext, scene: InteractiveScene):
scene.update(sim_dt)
# obtain quantities from simulation
ee_pose_w = robot.data.body_state_w[:, robot_entity_cfg.body_ids[0], 0:7]
ee_pose_w = robot.data.body_link_state_w[:, robot_entity_cfg.body_ids[0], 0:7]
# update marker positions
ee_marker.visualize(ee_pose_w[:, 0:3], ee_pose_w[:, 3:7])
goal_marker.visualize(ik_commands[:, 0:3] + scene.env_origins, ik_commands[:, 3:7])
......
......@@ -291,24 +291,26 @@ def update_states(
gravity = robot.root_physx_view.get_generalized_gravity_forces()[:, arm_joint_ids]
# Convert the Jacobian from world to root frame
jacobian_b = jacobian_w.clone()
root_rot_matrix = matrix_from_quat(quat_inv(robot.data.root_state_w[:, 3:7]))
root_rot_matrix = matrix_from_quat(quat_inv(robot.data.root_link_quat_w))
jacobian_b[:, :3, :] = torch.bmm(root_rot_matrix, jacobian_b[:, :3, :])
jacobian_b[:, 3:, :] = torch.bmm(root_rot_matrix, jacobian_b[:, 3:, :])
# Compute current pose of the end-effector
root_pose_w = robot.data.root_state_w[:, 0:7]
ee_pose_w = robot.data.body_state_w[:, ee_frame_idx, 0:7]
ee_pos_b, ee_quat_b = subtract_frame_transforms(
root_pose_w[:, 0:3], root_pose_w[:, 3:7], ee_pose_w[:, 0:3], ee_pose_w[:, 3:7]
)
root_pos_w = robot.data.root_link_pos_w
root_quat_w = robot.data.root_link_quat_w
ee_pos_w = robot.data.body_link_pos_w[:, ee_frame_idx]
ee_quat_w = robot.data.body_link_quat_w[:, ee_frame_idx]
ee_pos_b, ee_quat_b = subtract_frame_transforms(root_pos_w, root_quat_w, ee_pos_w, ee_quat_w)
root_pose_w = torch.cat([root_pos_w, root_quat_w], dim=-1)
ee_pose_w = torch.cat([ee_pos_w, ee_quat_w], dim=-1)
ee_pose_b = torch.cat([ee_pos_b, ee_quat_b], dim=-1)
# Compute the current velocity of the end-effector
ee_vel_w = robot.data.body_vel_w[:, ee_frame_idx, :] # Extract end-effector velocity in the world frame
root_vel_w = robot.data.root_vel_w # Extract root velocity in the world frame
ee_vel_w = robot.data.body_com_vel_w[:, ee_frame_idx, :] # Extract end-effector velocity in the world frame
root_vel_w = robot.data.root_com_vel_w # Extract root velocity in the world frame
relative_vel_w = ee_vel_w - root_vel_w # Compute the relative velocity in the world frame
ee_lin_vel_b = quat_rotate_inverse(robot.data.root_quat_w, relative_vel_w[:, 0:3]) # From world to root frame
ee_ang_vel_b = quat_rotate_inverse(robot.data.root_quat_w, relative_vel_w[:, 3:6])
ee_lin_vel_b = quat_rotate_inverse(robot.data.root_link_quat_w, relative_vel_w[:, 0:3]) # From world to root frame
ee_ang_vel_b = quat_rotate_inverse(robot.data.root_link_quat_w, relative_vel_w[:, 3:6])
ee_vel_b = torch.cat([ee_lin_vel_b, ee_ang_vel_b], dim=-1)
# Calculate the contact force
......
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