Skip to content

Extract teleop ROS publication builders #779

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
sgrizan-nv wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Extract teleop ROS publication builders #779

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 6 additions & 0 deletions examples/teleop_ros2/AGENTS.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,12 @@ SPDX-License-Identifier: Apache-2.0

- In source code files under this example, preserve the existing grouped/sorted organization for helpers, message builders, classes, and member functions: scan the surrounding order before inserting, and do not place helpers near call sites when the existing section is sorted.
- In Python integration test verifier code, do not use bare `assert` for runtime validation; Python optimization can disable it, so raise explicit exceptions from validators.
- Keep ROS publisher methods at one abstraction level: call one output-specific pure
builder, then publish or broadcast its result. Compose payload construction and
transport encoding inside the message-builder module, not in the node.
- Do not extract a one-use mapping helper merely to shorten a builder; keep the
complete output mapping together unless the helper is reused or represents an
independently tested contract.

## Rename consistency

Expand Down
288 changes: 202 additions & 86 deletions examples/teleop_ros2/python/messages.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,15 +2,18 @@
# All rights reserved.
# SPDX-License-Identifier: Apache-2.0

"""ROS message and msgpack payload builders for teleop_ros2_node."""
"""Pure publication-output builders for teleop_ros2_node."""

import time
from typing import Dict, Sequence

import msgpack
import msgpack_numpy as mnp
import numpy as np
from geometry_msgs.msg import Pose, PoseStamped
from geometry_msgs.msg import Pose, PoseStamped, TransformStamped, TwistStamped
from scipy.spatial.transform import Rotation
from sensor_msgs.msg import JointState
from std_msgs.msg import ByteMultiArray
from teleop_ros2_interfaces.msg import NamedPoseArray

from isaacteleop.retargeting_engine.interface import OptionalTensorGroup
Expand All @@ -26,6 +29,7 @@
from geometry import (
apply_manus_controller_to_hand_pose,
apply_transform_to_pose,
make_transform,
to_pose,
)
from tensor_group_helpers import (
Expand Down Expand Up @@ -100,17 +104,51 @@ def _compute_ee_pose_from_hand(
return pose


def _to_pose_stamped(pose: Pose, now, frame_id: str) -> PoseStamped:
msg = PoseStamped()
msg.header.stamp = now
msg.header.frame_id = frame_id
msg.pose = pose
def _to_msgpack_byte_multi_array(payload: Dict) -> ByteMultiArray:
encoded_payload = msgpack.packb(payload, default=mnp.encode)
msg = ByteMultiArray()
msg.data = tuple(bytes([value]) for value in encoded_payload)
return msg


def build_controller_payload(
left_ctrl: OptionalTensorGroup, right_ctrl: OptionalTensorGroup
) -> Dict:
def _wrist_tfs_from_ee_msg(
ee_msg: NamedPoseArray,
left_wrist_frame: str,
right_wrist_frame: str,
) -> list[TransformStamped]:
wrist_frames = {
"left": left_wrist_frame,
"right": right_wrist_frame,
}
transforms = []
for side, pose, is_valid in zip(ee_msg.name, ee_msg.pose, ee_msg.is_valid):
if not is_valid:
continue
transforms.append(
make_transform(
ee_msg.header.stamp,
ee_msg.header.frame_id,
wrist_frames[side],
[pose.position.x, pose.position.y, pose.position.z],
[
pose.orientation.x,
pose.orientation.y,
pose.orientation.z,
pose.orientation.w,
],
)
)
return transforms


def build_controller_msg(
left_ctrl: OptionalTensorGroup,
right_ctrl: OptionalTensorGroup,
) -> ByteMultiArray | None:
"""Build a msgpack controller message, or None when both inputs are absent."""
if left_ctrl.is_none and right_ctrl.is_none:
return None

def _as_list(ctrl, index):
if ctrl.is_none:
return [0.0, 0.0, 0.0]
Expand All @@ -126,73 +164,89 @@ def _as_float(ctrl, index):
return 0.0
return float(ctrl[index])

return {
"timestamp": time.time_ns(),
"left_thumbstick": [
_as_float(left_ctrl, ControllerInputIndex.THUMBSTICK_X),
_as_float(left_ctrl, ControllerInputIndex.THUMBSTICK_Y),
],
"right_thumbstick": [
_as_float(right_ctrl, ControllerInputIndex.THUMBSTICK_X),
_as_float(right_ctrl, ControllerInputIndex.THUMBSTICK_Y),
],
"left_trigger_value": _as_float(left_ctrl, ControllerInputIndex.TRIGGER_VALUE),
"right_trigger_value": _as_float(
right_ctrl, ControllerInputIndex.TRIGGER_VALUE
),
"left_squeeze_value": _as_float(left_ctrl, ControllerInputIndex.SQUEEZE_VALUE),
"right_squeeze_value": _as_float(
right_ctrl, ControllerInputIndex.SQUEEZE_VALUE
),
"left_aim_position": _as_list(left_ctrl, ControllerInputIndex.AIM_POSITION),
"right_aim_position": _as_list(right_ctrl, ControllerInputIndex.AIM_POSITION),
"left_grip_position": _as_list(left_ctrl, ControllerInputIndex.GRIP_POSITION),
"right_grip_position": _as_list(right_ctrl, ControllerInputIndex.GRIP_POSITION),
"left_aim_orientation": _as_quat(
left_ctrl, ControllerInputIndex.AIM_ORIENTATION
),
"right_aim_orientation": _as_quat(
right_ctrl, ControllerInputIndex.AIM_ORIENTATION
),
"left_grip_orientation": _as_quat(
left_ctrl, ControllerInputIndex.GRIP_ORIENTATION
),
"right_grip_orientation": _as_quat(
right_ctrl, ControllerInputIndex.GRIP_ORIENTATION
),
"left_primary_click": _as_float(left_ctrl, ControllerInputIndex.PRIMARY_CLICK),
"right_primary_click": _as_float(
right_ctrl, ControllerInputIndex.PRIMARY_CLICK
),
"left_secondary_click": _as_float(
left_ctrl, ControllerInputIndex.SECONDARY_CLICK
),
"right_secondary_click": _as_float(
right_ctrl, ControllerInputIndex.SECONDARY_CLICK
),
"left_thumbstick_click": _as_float(
left_ctrl, ControllerInputIndex.THUMBSTICK_CLICK
),
"right_thumbstick_click": _as_float(
right_ctrl, ControllerInputIndex.THUMBSTICK_CLICK
),
"left_menu_click": _as_float(left_ctrl, ControllerInputIndex.MENU_CLICK),
"right_menu_click": _as_float(right_ctrl, ControllerInputIndex.MENU_CLICK),
"left_is_active": not left_ctrl.is_none,
"right_is_active": not right_ctrl.is_none,
}
return _to_msgpack_byte_multi_array(
{
"timestamp": time.time_ns(),
"left_thumbstick": [
_as_float(left_ctrl, ControllerInputIndex.THUMBSTICK_X),
_as_float(left_ctrl, ControllerInputIndex.THUMBSTICK_Y),
],
"right_thumbstick": [
_as_float(right_ctrl, ControllerInputIndex.THUMBSTICK_X),
_as_float(right_ctrl, ControllerInputIndex.THUMBSTICK_Y),
],
"left_trigger_value": _as_float(
left_ctrl, ControllerInputIndex.TRIGGER_VALUE
),
"right_trigger_value": _as_float(
right_ctrl, ControllerInputIndex.TRIGGER_VALUE
),
"left_squeeze_value": _as_float(
left_ctrl, ControllerInputIndex.SQUEEZE_VALUE
),
"right_squeeze_value": _as_float(
right_ctrl, ControllerInputIndex.SQUEEZE_VALUE
),
"left_aim_position": _as_list(left_ctrl, ControllerInputIndex.AIM_POSITION),
"right_aim_position": _as_list(
right_ctrl, ControllerInputIndex.AIM_POSITION
),
"left_grip_position": _as_list(
left_ctrl, ControllerInputIndex.GRIP_POSITION
),
"right_grip_position": _as_list(
right_ctrl, ControllerInputIndex.GRIP_POSITION
),
"left_aim_orientation": _as_quat(
left_ctrl, ControllerInputIndex.AIM_ORIENTATION
),
"right_aim_orientation": _as_quat(
right_ctrl, ControllerInputIndex.AIM_ORIENTATION
),
"left_grip_orientation": _as_quat(
left_ctrl, ControllerInputIndex.GRIP_ORIENTATION
),
"right_grip_orientation": _as_quat(
right_ctrl, ControllerInputIndex.GRIP_ORIENTATION
),
"left_primary_click": _as_float(
left_ctrl, ControllerInputIndex.PRIMARY_CLICK
),
"right_primary_click": _as_float(
right_ctrl, ControllerInputIndex.PRIMARY_CLICK
),
"left_secondary_click": _as_float(
left_ctrl, ControllerInputIndex.SECONDARY_CLICK
),
"right_secondary_click": _as_float(
right_ctrl, ControllerInputIndex.SECONDARY_CLICK
),
"left_thumbstick_click": _as_float(
left_ctrl, ControllerInputIndex.THUMBSTICK_CLICK
),
"right_thumbstick_click": _as_float(
right_ctrl, ControllerInputIndex.THUMBSTICK_CLICK
),
"left_menu_click": _as_float(left_ctrl, ControllerInputIndex.MENU_CLICK),
"right_menu_click": _as_float(right_ctrl, ControllerInputIndex.MENU_CLICK),
"left_is_active": not left_ctrl.is_none,
"right_is_active": not right_ctrl.is_none,
}
)


def build_ee_msg_from_controllers(
def build_ee_output_from_controllers(
left_ctrl: OptionalTensorGroup,
right_ctrl: OptionalTensorGroup,
now,
frame_id: str,
left_wrist_frame: str,
right_wrist_frame: str,
transform_rot: Rotation | None = None,
transform_trans: Sequence[float] | None = None,
controller_uses_hands_source: bool = False,
) -> NamedPoseArray:
"""Build fixed left/right EE entries from controller aim poses."""
) -> tuple[NamedPoseArray, list[TransformStamped]]:
"""Build the controller-derived EE message and its valid wrist TFs."""
left_pose = _compute_ee_pose_from_controller(
left_ctrl,
"left",
Expand All @@ -207,18 +261,25 @@ def build_ee_msg_from_controllers(
transform_trans,
controller_uses_hands_source,
)
return _compose_ee_msg(left_pose, right_pose, now, frame_id)
ee_msg = _compose_ee_msg(left_pose, right_pose, now, frame_id)
return ee_msg, _wrist_tfs_from_ee_msg(
ee_msg,
left_wrist_frame,
right_wrist_frame,
)


def build_ee_msg_from_hands(
def build_ee_output_from_hands(
left_hand: OptionalTensorGroup,
right_hand: OptionalTensorGroup,
now,
frame_id: str,
left_wrist_frame: str,
right_wrist_frame: str,
transform_rot: Rotation | None = None,
transform_trans: Sequence[float] | None = None,
) -> NamedPoseArray:
"""Build fixed left/right EE entries from hand wrist poses."""
) -> tuple[NamedPoseArray, list[TransformStamped]]:
"""Build the hand-derived EE message and its valid wrist TFs."""
left_pose = _compute_ee_pose_from_hand(
left_hand,
transform_rot,
Expand All @@ -229,7 +290,12 @@ def build_ee_msg_from_hands(
transform_rot,
transform_trans,
)
return _compose_ee_msg(left_pose, right_pose, now, frame_id)
ee_msg = _compose_ee_msg(left_pose, right_pose, now, frame_id)
return ee_msg, _wrist_tfs_from_ee_msg(
ee_msg,
left_wrist_frame,
right_wrist_frame,
)


def build_finger_joints_msg(
Expand Down Expand Up @@ -267,18 +333,25 @@ def build_finger_joints_msg(
return finger_joints_msg


def build_full_body_payload(full_body: OptionalTensorGroup) -> Dict:
def build_full_body_msg(
full_body: OptionalTensorGroup,
) -> ByteMultiArray | None:
"""Build a msgpack full-body message, or None when the input is absent."""
if full_body.is_none:
return None

positions = np.asarray(full_body[FullBodyInputIndex.JOINT_POSITIONS])
orientations = np.asarray(full_body[FullBodyInputIndex.JOINT_ORIENTATIONS])
valid = np.asarray(full_body[FullBodyInputIndex.JOINT_VALID])

return {
"timestamp": time.time_ns(),
"joint_names": BODY_JOINT_NAMES,
"joint_positions": [[float(v) for v in pos] for pos in positions],
"joint_orientations": [[float(v) for v in ori] for ori in orientations],
"joint_valid": [bool(v) for v in valid],
}
return _to_msgpack_byte_multi_array(
{
"timestamp": time.time_ns(),
"joint_names": BODY_JOINT_NAMES,
"joint_positions": [[float(v) for v in pos] for pos in positions],
"joint_orientations": [[float(v) for v in ori] for ori in orientations],
"joint_valid": [bool(v) for v in valid],
}
)


def build_hand_msg(
Expand Down Expand Up @@ -318,14 +391,15 @@ def build_hand_msg(
return msg


def build_head_msg(
def build_head_output(
head: OptionalTensorGroup,
now,
frame_id: str,
child_frame: str,
transform_rot: Rotation | None = None,
transform_trans: Sequence[float] | None = None,
) -> PoseStamped | None:
"""Build a PoseStamped for the head pose, or None when head is invalid."""
) -> tuple[PoseStamped, TransformStamped] | None:
"""Build the head pose message and matching TF, or None when invalid."""
if not head_is_valid(head):
return None

Expand All @@ -335,4 +409,46 @@ def build_head_msg(
if transform_rot is not None or transform_trans is not None:
pose = apply_transform_to_pose(pose, transform_rot, transform_trans)

return _to_pose_stamped(pose, now, frame_id)
head_msg = PoseStamped()
head_msg.header.stamp = now
head_msg.header.frame_id = frame_id
head_msg.pose = pose
head_tf = make_transform(
now,
frame_id,
child_frame,
[pose.position.x, pose.position.y, pose.position.z],
[
pose.orientation.x,
pose.orientation.y,
pose.orientation.z,
pose.orientation.w,
],
)
return head_msg, head_tf


def build_root_command_output(
root_command: OptionalTensorGroup,
now,
frame_id: str,
) -> tuple[TwistStamped, PoseStamped] | None:
"""Build root twist and height-pose messages, or None for an absent command."""
if root_command.is_none:
return None

command = np.asarray(root_command[0])
twist_msg = TwistStamped()
twist_msg.header.stamp = now
twist_msg.header.frame_id = frame_id
twist_msg.twist.linear.x = float(command[0])
twist_msg.twist.linear.y = float(command[1])
twist_msg.twist.linear.z = 0.0
twist_msg.twist.angular.z = float(command[2])

pose_msg = PoseStamped()
pose_msg.header.stamp = now
pose_msg.header.frame_id = frame_id
pose_msg.pose.position.z = float(command[3])
pose_msg.pose.orientation.w = 1.0
return twist_msg, pose_msg
Loading
Loading