diff --git a/examples/teleop_ros2/AGENTS.md b/examples/teleop_ros2/AGENTS.md index 5b0431dca..3c0682d98 100644 --- a/examples/teleop_ros2/AGENTS.md +++ b/examples/teleop_ros2/AGENTS.md @@ -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 diff --git a/examples/teleop_ros2/python/messages.py b/examples/teleop_ros2/python/messages.py index d52b68b4b..a586f1ba0 100644 --- a/examples/teleop_ros2/python/messages.py +++ b/examples/teleop_ros2/python/messages.py @@ -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 @@ -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 ( @@ -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] @@ -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", @@ -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, @@ -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( @@ -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( @@ -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 @@ -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 diff --git a/examples/teleop_ros2/python/teleop_ros2_node.py b/examples/teleop_ros2/python/teleop_ros2_node.py index 1c498d317..ca42d9693 100755 --- a/examples/teleop_ros2/python/teleop_ros2_node.py +++ b/examples/teleop_ros2/python/teleop_ros2_node.py @@ -39,16 +39,10 @@ """ import time -from typing import List -import msgpack -import msgpack_numpy as mnp -import numpy as np import rclpy from geometry_msgs.msg import ( - Pose, PoseStamped, - TransformStamped, TwistStamped, ) from rclpy.node import Node @@ -59,15 +53,15 @@ from isaacteleop.cloudxr import CloudXRLauncher from isaacteleop.teleop_session_manager import SessionMode, TeleopSession -from geometry import make_transform from messages import ( - build_controller_payload, - build_ee_msg_from_controllers, - build_ee_msg_from_hands, + build_controller_msg, + build_ee_output_from_controllers, + build_ee_output_from_hands, build_finger_joints_msg, - build_full_body_payload, + build_full_body_msg, build_hand_msg, - build_head_msg, + build_head_output, + build_root_command_output, ) from teleop_profiles import ( PublishType, @@ -80,9 +74,6 @@ create_node_parameters, ) from session_config import build_session_config -from tensor_group_helpers import ( - head_is_valid, -) class TeleopRos2Node(Node): @@ -98,39 +89,6 @@ def __init__(self) -> None: self._create_publishers() self._config = build_session_config(self._params) - def _build_wrist_tfs( - self, - ee_msg: NamedPoseArray, - ) -> List[TransformStamped]: - """Build wrist TF transforms from valid EE pose entries.""" - tfs = [] - wrist_frames = { - "left": self._params.left_wrist_frame, - "right": self._params.right_wrist_frame, - } - - def _get_orientation(pose: Pose) -> List[float]: - return [ - pose.orientation.x, - pose.orientation.y, - pose.orientation.z, - pose.orientation.w, - ] - - for side, pose, is_valid in zip(ee_msg.name, ee_msg.pose, ee_msg.is_valid): - if not is_valid: - continue - tfs.append( - make_transform( - ee_msg.header.stamp, - ee_msg.header.frame_id, - wrist_frames[side], - [pose.position.x, pose.position.y, pose.position.z], - _get_orientation(pose), - ) - ) - return tfs - def _create_publishers(self) -> None: self._pub_hand = self.create_publisher(NamedPoseArray, "xr_teleop/hand", 10) self._pub_ee_pose = self.create_publisher( @@ -154,33 +112,28 @@ def _create_publishers(self) -> None: self._pub_head = self.create_publisher(PoseStamped, "xr_teleop/head_pose", 10) def _publish_ee_pose_from_controllers(self, result: SessionResult, now) -> None: - left_ctrl = result["controller_left"] - right_ctrl = result["controller_right"] - ee_pose_msg = build_ee_msg_from_controllers( - left_ctrl, - right_ctrl, + ee_pose_msg, wrist_tfs = build_ee_output_from_controllers( + result["controller_left"], + result["controller_right"], now, self._params.world_frame, + self._params.left_wrist_frame, + self._params.right_wrist_frame, self._params.transform_rotation, self._params.transform_translation, self._params.controller_uses_hands_source, ) self._pub_ee_pose.publish(ee_pose_msg) - wrist_tfs = self._build_wrist_tfs(ee_pose_msg) if wrist_tfs: self._tf_broadcaster.sendTransform(wrist_tfs) def _publish_controller_payload(self, result: SessionResult) -> None: - left_ctrl = result["controller_left"] - right_ctrl = result["controller_right"] - if left_ctrl.is_none and right_ctrl.is_none: - return - - controller_payload = build_controller_payload(left_ctrl, right_ctrl) - payload = msgpack.packb(controller_payload, default=mnp.encode) - controller_msg = ByteMultiArray() - controller_msg.data = tuple(bytes([a]) for a in payload) - self._pub_controller.publish(controller_msg) + controller_msg = build_controller_msg( + result["controller_left"], + result["controller_right"], + ) + if controller_msg is not None: + self._pub_controller.publish(controller_msg) def _publish_finger_joints(self, result: SessionResult, now) -> None: finger_joints_msg = build_finger_joints_msg( @@ -193,15 +146,9 @@ def _publish_finger_joints(self, result: SessionResult, now) -> None: self._pub_finger_joints.publish(finger_joints_msg) def _publish_full_body_payload(self, result: SessionResult) -> None: - full_body_data = result["full_body"] - if full_body_data.is_none: - return - - body_payload = build_full_body_payload(full_body_data) - payload = msgpack.packb(body_payload, default=mnp.encode) - body_msg = ByteMultiArray() - body_msg.data = tuple(bytes([a]) for a in payload) - self._pub_full_body.publish(body_msg) + body_msg = build_full_body_msg(result["full_body"]) + if body_msg is not None: + self._pub_full_body.publish(body_msg) def _publish_hand_poses(self, result: SessionResult, now) -> None: hand_msg = build_hand_msg( @@ -215,72 +162,47 @@ def _publish_hand_poses(self, result: SessionResult, now) -> None: self._pub_hand.publish(hand_msg) def _publish_ee_pose_from_hands(self, result: SessionResult, now) -> None: - left_hand = result["hand_left"] - right_hand = result["hand_right"] - ee_pose_msg = build_ee_msg_from_hands( - left_hand, - right_hand, + ee_pose_msg, wrist_tfs = build_ee_output_from_hands( + result["hand_left"], + result["hand_right"], now, self._params.world_frame, + self._params.left_wrist_frame, + self._params.right_wrist_frame, self._params.transform_rotation, self._params.transform_translation, ) self._pub_ee_pose.publish(ee_pose_msg) - wrist_tfs = self._build_wrist_tfs(ee_pose_msg) if wrist_tfs: self._tf_broadcaster.sendTransform(wrist_tfs) def _publish_head(self, result: SessionResult, now) -> None: - head = result["head"] - if not head_is_valid(head): - return - - head_msg = build_head_msg( - head, + head_output = build_head_output( + result["head"], now, self._params.world_frame, + self._params.head_frame, self._params.transform_rotation, self._params.transform_translation, ) - if head_msg is None: + if head_output is None: return + head_msg, head_tf = head_output self._pub_head.publish(head_msg) - pose = head_msg.pose - head_tf = make_transform( - now, - self._params.world_frame, - self._params.head_frame, - [pose.position.x, pose.position.y, pose.position.z], - [ - pose.orientation.x, - pose.orientation.y, - pose.orientation.z, - pose.orientation.w, - ], - ) self._tf_broadcaster.sendTransform(head_tf) def _publish_root_command(self, result: SessionResult, now) -> None: - root_command = result["root_command"] - if root_command.is_none: + root_output = build_root_command_output( + result["root_command"], + now, + self._params.world_frame, + ) + if root_output is None: return - cmd = np.asarray(root_command[0]) - twist_msg = TwistStamped() - twist_msg.header.stamp = now - twist_msg.header.frame_id = self._params.world_frame - twist_msg.twist.linear.x = float(cmd[0]) - twist_msg.twist.linear.y = float(cmd[1]) - twist_msg.twist.linear.z = 0.0 - twist_msg.twist.angular.z = float(cmd[2]) + twist_msg, pose_msg = root_output self._pub_root_twist.publish(twist_msg) - - pose_msg = PoseStamped() - pose_msg.header.stamp = now - pose_msg.header.frame_id = self._params.world_frame - pose_msg.pose.position.z = float(cmd[3]) - pose_msg.pose.orientation.w = 1.0 self._pub_root_pose.publish(pose_msg) def _run_session_loop(self, launcher: CloudXRLauncher | None = None) -> int: diff --git a/examples/teleop_ros2/python/tests/test_messages.py b/examples/teleop_ros2/python/tests/test_messages.py index 6a96211b6..169265ee0 100644 --- a/examples/teleop_ros2/python/tests/test_messages.py +++ b/examples/teleop_ros2/python/tests/test_messages.py @@ -2,25 +2,45 @@ # All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Tests for tensor-group to ROS-message conversion.""" +"""Behavioral tests for teleoperation publication-output builders.""" +import msgpack import numpy as np from builtin_interfaces.msg import Time +from std_msgs.msg import ByteMultiArray -from isaacteleop.retargeting_engine.interface import OptionalTensorGroup, TensorGroup +from isaacteleop.retargeting_engine.interface import ( + OptionalTensorGroup, + TensorGroup, + TensorGroupType, +) from isaacteleop.retargeting_engine.tensor_types import ( + DLDataType, + NDArrayType, + NUM_BODY_JOINTS_PICO, NUM_HAND_JOINTS, ControllerInput, ControllerInputIndex, + FullBodyInput, + FullBodyInputIndex, HandInput, HandInputIndex, + HandJointIndex, + HeadPose, + HeadPoseIndex, + RobotHandJoints, ) -from constants import HAND_POSE_NAMES +from constants import BODY_JOINT_NAMES, HAND_POSE_NAMES from messages import ( - build_controller_payload, - build_ee_msg_from_controllers, + build_controller_msg, + build_ee_output_from_controllers, + build_ee_output_from_hands, + build_finger_joints_msg, + build_full_body_msg, build_hand_msg, + build_head_output, + build_root_command_output, ) @@ -51,6 +71,21 @@ def _active_controller() -> TensorGroup: return controller +def _active_full_body() -> TensorGroup: + full_body = TensorGroup(FullBodyInput()) + positions = np.zeros((NUM_BODY_JOINTS_PICO, 3), dtype=np.float32) + positions[0] = [1.0, 2.0, 3.0] + orientations = np.zeros((NUM_BODY_JOINTS_PICO, 4), dtype=np.float32) + orientations[:, 3] = 1.0 + valid = np.zeros(NUM_BODY_JOINTS_PICO, dtype=np.uint8) + valid[0] = 1 + + full_body[FullBodyInputIndex.JOINT_POSITIONS] = positions + full_body[FullBodyInputIndex.JOINT_ORIENTATIONS] = orientations + full_body[FullBodyInputIndex.JOINT_VALID] = valid + return full_body + + def _active_hand() -> TensorGroup: hand = TensorGroup(HandInput()) positions = np.zeros((NUM_HAND_JOINTS, 3), dtype=np.float32) @@ -66,40 +101,43 @@ def _active_hand() -> TensorGroup: return hand -def test_ee_message_has_fixed_entries_and_invalid_placeholder() -> None: - left = _active_controller() - right = OptionalTensorGroup(ControllerInput()) - - msg = build_ee_msg_from_controllers(left, right, Time(), "world") +def _active_head() -> TensorGroup: + head = TensorGroup(HeadPose()) + head[HeadPoseIndex.POSITION] = np.array([1.0, 2.0, 3.0], dtype=np.float32) + head[HeadPoseIndex.ORIENTATION] = np.array([0.0, 0.0, 0.0, 1.0], dtype=np.float32) + head[HeadPoseIndex.IS_VALID] = True + return head - assert list(msg.name) == ["left", "right"] - assert list(msg.is_valid) == [True, False] - assert msg.pose[0].position.x == 4.0 - assert msg.pose[1].position.x == 0.0 - assert msg.pose[1].orientation.w == 1.0 +def _decode_payload(msg: ByteMultiArray) -> dict: + return msgpack.unpackb(b"".join(msg.data), raw=False) -def test_hand_message_has_stable_names_and_absent_side_placeholders() -> None: - left = _active_hand() - right = OptionalTensorGroup(HandInput()) - msg = build_hand_msg(left, right, Time(), "world") - - expected_names = [ - f"{side}_{name}" for side in ("left", "right") for name in HAND_POSE_NAMES - ] - assert list(msg.name) == expected_names - assert len(msg.pose) == len(expected_names) - assert all(msg.is_valid[: len(HAND_POSE_NAMES)]) - assert not any(msg.is_valid[len(HAND_POSE_NAMES) :]) - assert "left_PALM" not in msg.name +def _root_command(values: list[float] | None) -> OptionalTensorGroup: + group_type = TensorGroupType( + "root_command", + [ + NDArrayType( + "command", + shape=(4,), + dtype=DLDataType.FLOAT, + dtype_bits=32, + ) + ], + ) + root_command = OptionalTensorGroup(group_type) + if values is not None: + root_command[0] = np.asarray(values, dtype=np.float32) + return root_command -def test_controller_payload_schema_and_absent_side_defaults() -> None: - payload = build_controller_payload( +def test_build_controller_msg_encodes_schema_and_absent_side_defaults() -> None: + msg = build_controller_msg( _active_controller(), OptionalTensorGroup(ControllerInput()) ) + assert msg is not None + payload = _decode_payload(msg) assert set(payload) == { "timestamp", "left_thumbstick", @@ -133,3 +171,208 @@ def test_controller_payload_schema_and_absent_side_defaults() -> None: assert payload["left_thumbstick"] == [0.25, -0.5] assert payload["right_aim_position"] == [0.0, 0.0, 0.0] assert payload["right_aim_orientation"] == [0.0, 0.0, 0.0, 1.0] + + +def test_build_controller_msg_returns_none_when_both_sides_absent() -> None: + assert ( + build_controller_msg( + OptionalTensorGroup(ControllerInput()), + OptionalTensorGroup(ControllerInput()), + ) + is None + ) + + +def test_build_ee_output_from_controllers_keeps_message_and_tfs_consistent() -> None: + msg, transforms = build_ee_output_from_controllers( + _active_controller(), + OptionalTensorGroup(ControllerInput()), + Time(sec=12, nanosec=34), + "world", + "left_wrist", + "right_wrist", + ) + + assert msg.header.stamp.sec == 12 + assert msg.header.stamp.nanosec == 34 + assert msg.header.frame_id == "world" + assert list(msg.name) == ["left", "right"] + assert list(msg.is_valid) == [True, False] + assert msg.pose[0].position.x == 4.0 + assert msg.pose[1].position.x == 0.0 + assert msg.pose[1].orientation.w == 1.0 + assert len(transforms) == 1 + transform = transforms[0] + assert transform.header.stamp.sec == 12 + assert transform.header.stamp.nanosec == 34 + assert transform.header.frame_id == "world" + assert transform.child_frame_id == "left_wrist" + assert transform.transform.translation.x == 4.0 + assert transform.transform.translation.y == 5.0 + assert transform.transform.translation.z == 6.0 + assert transform.transform.rotation.w == 1.0 + + +def test_build_ee_output_from_hands_uses_valid_wrist_entries() -> None: + msg, transforms = build_ee_output_from_hands( + _active_hand(), + OptionalTensorGroup(HandInput()), + Time(), + "world", + "left_wrist", + "right_wrist", + ) + + wrist_x = float(HandJointIndex.WRIST) + assert list(msg.name) == ["left", "right"] + assert list(msg.is_valid) == [True, False] + assert msg.pose[0].position.x == wrist_x + assert len(transforms) == 1 + assert transforms[0].child_frame_id == "left_wrist" + assert transforms[0].transform.translation.x == wrist_x + + +def test_build_finger_joints_msg_combines_present_sides() -> None: + left = TensorGroup( + RobotHandJoints("left_finger_joints", ["left_index", "left_thumb"]) + ) + left[0] = 1.0 + left[1] = 2.0 + right_type = RobotHandJoints("right_finger_joints", ["right_index"]) + + msg = build_finger_joints_msg( + left, + OptionalTensorGroup(right_type), + Time(sec=12, nanosec=34), + "world", + ) + + assert msg is not None + assert msg.header.stamp.sec == 12 + assert msg.header.stamp.nanosec == 34 + assert msg.header.frame_id == "world" + assert list(msg.name) == ["left_index", "left_thumb"] + assert list(msg.position) == [1.0, 2.0] + + +def test_build_finger_joints_msg_returns_none_when_both_sides_are_absent() -> None: + left_type = RobotHandJoints("left_finger_joints", ["left_index"]) + right_type = RobotHandJoints("right_finger_joints", ["right_index"]) + + assert ( + build_finger_joints_msg( + OptionalTensorGroup(left_type), + OptionalTensorGroup(right_type), + Time(), + "world", + ) + is None + ) + + +def test_build_full_body_msg_encodes_schema() -> None: + msg = build_full_body_msg(_active_full_body()) + + assert msg is not None + payload = _decode_payload(msg) + assert set(payload) == { + "timestamp", + "joint_names", + "joint_positions", + "joint_orientations", + "joint_valid", + } + assert isinstance(payload["timestamp"], int) + assert payload["joint_names"] == BODY_JOINT_NAMES + assert payload["joint_positions"][0] == [1.0, 2.0, 3.0] + assert payload["joint_orientations"][0] == [0.0, 0.0, 0.0, 1.0] + assert payload["joint_valid"][0] is True + assert not any(payload["joint_valid"][1:]) + + +def test_build_full_body_msg_returns_none_when_input_is_absent() -> None: + assert build_full_body_msg(OptionalTensorGroup(FullBodyInput())) is None + + +def test_build_hand_msg_keeps_stable_names_and_absent_side_placeholders() -> None: + msg = build_hand_msg( + _active_hand(), + OptionalTensorGroup(HandInput()), + Time(), + "world", + ) + + expected_names = [ + f"{side}_{name}" for side in ("left", "right") for name in HAND_POSE_NAMES + ] + assert list(msg.name) == expected_names + assert len(msg.pose) == len(expected_names) + assert all(msg.is_valid[: len(HAND_POSE_NAMES)]) + assert not any(msg.is_valid[len(HAND_POSE_NAMES) :]) + assert "left_PALM" not in msg.name + + +def test_build_head_output_keeps_pose_and_tf_consistent() -> None: + output = build_head_output( + _active_head(), + Time(sec=12, nanosec=34), + "world", + "head", + ) + + assert output is not None + head_msg, transform = output + assert head_msg.header.stamp.sec == 12 + assert head_msg.header.stamp.nanosec == 34 + assert head_msg.header.frame_id == "world" + assert head_msg.pose.position.x == 1.0 + assert head_msg.pose.position.y == 2.0 + assert head_msg.pose.position.z == 3.0 + assert head_msg.pose.orientation.w == 1.0 + assert transform.header.stamp.sec == 12 + assert transform.header.stamp.nanosec == 34 + assert transform.header.frame_id == "world" + assert transform.child_frame_id == "head" + assert transform.transform.translation.x == 1.0 + assert transform.transform.translation.y == 2.0 + assert transform.transform.translation.z == 3.0 + assert transform.transform.rotation.w == 1.0 + + +def test_build_head_output_returns_none_when_head_is_absent() -> None: + assert ( + build_head_output( + OptionalTensorGroup(HeadPose()), + Time(), + "world", + "head", + ) + is None + ) + + +def test_build_root_command_output_maps_twist_and_height() -> None: + output = build_root_command_output( + _root_command([1.0, -2.0, 3.0, 0.75]), + Time(sec=12, nanosec=34), + "world", + ) + + assert output is not None + twist_msg, pose_msg = output + assert twist_msg.header.stamp.sec == 12 + assert twist_msg.header.stamp.nanosec == 34 + assert twist_msg.header.frame_id == "world" + assert twist_msg.twist.linear.x == 1.0 + assert twist_msg.twist.linear.y == -2.0 + assert twist_msg.twist.linear.z == 0.0 + assert twist_msg.twist.angular.z == 3.0 + assert pose_msg.header.stamp.sec == 12 + assert pose_msg.header.stamp.nanosec == 34 + assert pose_msg.header.frame_id == "world" + assert pose_msg.pose.position.z == 0.75 + assert pose_msg.pose.orientation.w == 1.0 + + +def test_build_root_command_output_returns_none_when_input_is_absent() -> None: + assert build_root_command_output(_root_command(None), Time(), "world") is None