feat: add head control and joint data decoder

tg3_local_teleop/tg3_local_teleop.py

@@ -27,7 +27,7 @@ from brainco_hand_msgs.msg import MotorStatus, SetMotorMulti
 from diagnostic_msgs.msg import DiagnosticStatus
 from geometry_msgs.msg import TwistStamped
 from rclpy.node import Node
-from ros2_bridge_msgs.msg import ArmStatus
+from ros2_bridge_msgs.msg import ArmStatus, HeadCtrl, MotorCtrl, RobotState
 from sensor_msgs.msg import JointState
 from std_msgs.msg import String
 from std_srvs.srv import Trigger
@@ -41,6 +41,7 @@ from gesture_toggle import (
     right_b_pressed,
     select_hand_target,
 )
+from head_control import HeadPitchController, head_pitch_command_axis, left_z_pressed
 
 
 JOINT_NAMES = [
@@ -453,6 +454,25 @@ class LocalTeleopBridge(Node):
         self.locomotion_enabled = bool(
             self.locomotion_cfg.get("enabled", False)
         )
+        self.head_cfg = config.get("head", {})
+        self.head_enabled = bool(self.head_cfg.get("enabled", False))
+        self.head_pitch_controller = HeadPitchController(
+            deadzone=float(self.head_cfg.get("joystick_deadzone", 0.1)),
+            max_speed_rad_s=float(
+                self.head_cfg.get("max_pitch_speed_rad_s", 0.2)
+            ),
+            max_accel_rad_s2=float(
+                self.head_cfg.get("max_pitch_accel_rad_s2", 0.5)
+            ),
+            min_pitch_rad=float(
+                self.head_cfg.get("min_pitch_rad", -0.1745329252)
+            ),
+            max_pitch_rad=float(
+                self.head_cfg.get("max_pitch_rad", 0.872664626)
+            ),
+            axis_sign=float(self.head_cfg.get("pitch_axis_sign", -1.0)),
+            max_dt_s=float(self.head_cfg.get("max_integration_dt_s", 0.1)),
+        )
         self.data_collection_cfg = config.get("data_collection", {})
         self.data_collection_enabled = bool(
             self.data_collection_cfg.get("enabled", False)
@@ -504,6 +524,17 @@ class LocalTeleopBridge(Node):
         self.walk_zero_frames_remaining = 0
         self.walk_require_neutral = True
         self.last_walk_publish_at = 0.0
+        self.head_publisher = None
+        self.robot_head_positions: dict[int, float] = {}
+        self.robot_head_errors: dict[int, int] = {}
+        self.robot_head_at = 0.0
+        self.expected_head_messages: list[tuple[tuple[Any, ...], float]] = []
+        self.foreign_head_command_at = 0.0
+        self.foreign_head_command_count = 0
+        self.head_publish_count = 0
+        self.last_head_publish_at = 0.0
+        self.head_output_reasons: list[str] = []
+        self.head_z_pressed: bool | None = None
 
         self.publisher = self.create_publisher(JointState, ros_cfg["command_topic"], 10)
         self.walk_publisher = None
@@ -511,6 +542,22 @@ class LocalTeleopBridge(Node):
             self.walk_publisher = self.create_publisher(
                 TwistStamped, self.locomotion_cfg["command_topic"], 10
             )
+        if self.head_enabled:
+            self.head_publisher = self.create_publisher(
+                HeadCtrl, self.head_cfg["command_topic"], 10
+            )
+            self.create_subscription(
+                RobotState,
+                self.head_cfg["state_topic"],
+                self._on_robot_state,
+                10,
+            )
+            self.create_subscription(
+                HeadCtrl,
+                self.head_cfg["command_topic"],
+                self._on_head_command,
+                10,
+            )
         self.create_subscription(
             DiagnosticStatus, ros_cfg["rl_state_topic"], self._on_rl_state, 10
         )
@@ -597,6 +644,7 @@ class LocalTeleopBridge(Node):
             f"local bridge started in {mode}; source={self.source.description}; "
             f"target={ros_cfg['command_topic']}; brainco_hands={self.hands_enabled}; "
             f"locomotion={self.locomotion_enabled}; "
+            f"head_pitch={self.head_enabled}; "
             f"data_collection={self.data_collection_enabled}"
         )
 
@@ -620,6 +668,20 @@ class LocalTeleopBridge(Node):
         self.robot_arm_errors = [int(by_id[motor_id].error) for motor_id in MOTOR_IDS]
         self.robot_arm_at = time.monotonic()
 
+    def _on_robot_state(self, msg: RobotState) -> None:
+        by_id = {int(motor.name): motor for motor in msg.head.status}
+        pitch_id = int(self.head_cfg.get("pitch_motor_id", 2))
+        if pitch_id not in by_id:
+            return
+        positions = {motor_id: float(motor.pos) for motor_id, motor in by_id.items()}
+        if not all(math.isfinite(value) for value in positions.values()):
+            return
+        self.robot_head_positions = positions
+        self.robot_head_errors = {
+            motor_id: int(motor.error) for motor_id, motor in by_id.items()
+        }
+        self.robot_head_at = time.monotonic()
+
     def _on_hand_status(self, side: str, msg: MotorStatus) -> None:
         positions = [int(value) for value in msg.positions]
         states = [int(value) for value in msg.states]
@@ -899,6 +961,45 @@ class LocalTeleopBridge(Node):
                 "foreign BrainCo hand command source detected", success=False
             )
 
+    @staticmethod
+    def _head_message_signature(msg: HeadCtrl) -> tuple[Any, ...]:
+        return (
+            int(msg.mode),
+            int(msg.label),
+            int(msg.reserved),
+            tuple(
+                (
+                    int(motor.name),
+                    float(motor.kp),
+                    float(motor.kd),
+                    float(motor.pos),
+                    float(motor.spd),
+                    float(motor.tor),
+                    float(motor.cur),
+                    str(motor.joint_ids),
+                )
+                for motor in msg.ctrl
+            ),
+        )
+
+    def _on_head_command(self, msg: HeadCtrl) -> None:
+        now = time.monotonic()
+        window = float(self.head_cfg.get("self_command_window_s", 0.25))
+        self.expected_head_messages = [
+            expected
+            for expected in self.expected_head_messages
+            if now - expected[1] <= window
+        ]
+        signature = self._head_message_signature(msg)
+        if any(signature == expected[0] for expected in self.expected_head_messages):
+            return
+        self.foreign_head_command_at = now
+        self.foreign_head_command_count += 1
+        if self.head_pitch_controller.active:
+            self.get_logger().warning(
+                "external /head/cmd detected; local head-pitch output paused"
+            )
+
     def _on_command_topic(self, msg: JointState) -> None:
         if msg.header.frame_id != self.FRAME_ID and len(msg.position) in (14, 16):
             if not self.foreign_source_seen:
@@ -1001,6 +1102,8 @@ class LocalTeleopBridge(Node):
         # decision, so a matching STOP or safety teardown always wins over an
         # L3 press observed in the same control tick.
         self._tick_data_collection(now, sample)
+        if self.head_enabled:
+            self._tick_head(now, sample)
 
         if self.right_point_gesture_enabled:
             gesture_toggled = self.right_point_gesture.update(
@@ -1288,6 +1391,8 @@ class LocalTeleopBridge(Node):
             self.right_point_gesture.new_session()
         if self.right_a_pose_enabled:
             self.right_a_pose.new_session()
+        if self.head_enabled:
+            self.head_pitch_controller.new_session()
         if getattr(self, "data_collection_enabled", False):
             self.data_capture_session_id = session_id
             self.data_collection_gate.new_session()
@@ -1360,6 +1465,8 @@ class LocalTeleopBridge(Node):
             self.right_point_gesture.new_session()
         if self.right_a_pose_enabled:
             self.right_a_pose.new_session()
+        if self.head_enabled:
+            self.head_pitch_controller.new_session()
         if getattr(self, "data_collection_enabled", False):
             self.data_capture_session_id = "direct_" + uuid.uuid4().hex
             self.data_collection_gate.new_session()
@@ -1398,6 +1505,8 @@ class LocalTeleopBridge(Node):
             self.right_point_gesture.disarm(reason)
         if self.right_a_pose_enabled:
             self.right_a_pose.disarm(reason)
+        if self.head_enabled:
+            self.head_pitch_controller.disarm(reason)
         self.armed = False
         self.hand_output_ready = False
         self.runtime_hand_output_reasons = []
@@ -1438,6 +1547,9 @@ class LocalTeleopBridge(Node):
             # convention requested for this installation.
             raw_forward = float(left_joystick[0])
             raw_yaw = float(right_joystick[1])
+            head_mode = left_z_pressed(sample.data)
+            if head_mode is None:
+                raise ValueError("left Z state is malformed")
             deadzone = float(cfg["joystick_deadzone"])
             forward_expo = float(cfg["joystick_expo"])
             yaw_expo = float(cfg.get("yaw_joystick_expo", forward_expo))
@@ -1447,6 +1559,11 @@ class LocalTeleopBridge(Node):
             shaped_yaw = self._shape_joystick_axis(
                 raw_yaw, deadzone, yaw_expo
             )
+            # Left Z reserves the entire right stick for head-pitch mode.
+            # Suppress yaw so a small sideways component while looking up or
+            # down cannot turn the robot.
+            if head_mode:
+                shaped_yaw = 0.0
             signed_forward = shaped_forward * float(
                 cfg.get("forward_axis_sign", 1.0)
             )
@@ -1533,6 +1650,96 @@ class LocalTeleopBridge(Node):
         self.last_walk_publish_at = now
         self.walk_publish_count += 1
 
+    def _tick_head(self, now: float, sample: ArmSnapshot | None) -> None:
+        pitch_id = int(self.head_cfg.get("pitch_motor_id", 2))
+        pitch_position = self.robot_head_positions.get(pitch_id)
+        feedback_timeout = float(self.head_cfg.get("feedback_timeout_s", 0.25))
+        feedback_age = (
+            math.inf if self.robot_head_at == 0.0 else now - self.robot_head_at
+        )
+        feedback_healthy = (
+            pitch_position is not None
+            and math.isfinite(float(pitch_position))
+            and feedback_age <= feedback_timeout
+            and int(self.robot_head_errors.get(pitch_id, -1)) == 0
+        )
+        input_timeout = float(self.head_cfg.get("input_timeout_s", 0.3))
+        input_healthy = (
+            self.armed
+            and sample is not None
+            and now - sample.received_at <= input_timeout
+        )
+        self.head_z_pressed = (
+            None if sample is None else left_z_pressed(sample.data)
+        )
+        axis = None if sample is None else head_pitch_command_axis(sample.data)
+        quiet_s = float(self.head_cfg.get("foreign_command_quiet_s", 1.0))
+        foreign_age = (
+            math.inf
+            if self.foreign_head_command_at == 0.0
+            else now - self.foreign_head_command_at
+        )
+        external_busy = foreign_age <= quiet_s
+
+        reasons: list[str] = []
+        if not self.armed:
+            reasons.append("teleoperation is not armed")
+        if not input_healthy:
+            reasons.append("right-stick input is unavailable/stale")
+        if axis is None:
+            reasons.append("left Z or right-stick vertical input is malformed")
+        elif not self.head_z_pressed:
+            reasons.append("left Z is not held; head-pitch control is idle")
+        if not feedback_healthy:
+            reasons.append("head pitch feedback is unavailable/stale or in error")
+        if external_busy:
+            reasons.append("external /head/cmd publisher has an active lease")
+        self.head_output_reasons = reasons
+
+        target = self.head_pitch_controller.update(
+            now,
+            armed=self.armed,
+            input_healthy=input_healthy,
+            feedback_healthy=feedback_healthy,
+            external_busy=external_busy,
+            axis=axis,
+            feedback_position_rad=pitch_position,
+        )
+        if target is not None:
+            self._publish_head_pitch(target, now)
+
+    def _publish_head_pitch(self, target_rad: float, now: float) -> None:
+        if not self.allow_publish or self.head_publisher is None:
+            return
+        motor = MotorCtrl()
+        motor.name = int(self.head_cfg.get("pitch_motor_id", 2))
+        motor.kp = 0.0
+        motor.kd = 0.0
+        motor.pos = float(target_rad)
+        motor.spd = float(self.head_cfg.get("command_speed_rad_s", 0.2))
+        motor.tor = 0.0
+        motor.cur = float(self.head_cfg.get("max_current_a", 1.0))
+        motor.joint_ids = ""
+
+        msg = HeadCtrl()
+        msg.header.stamp = self.get_clock().now().to_msg()
+        msg.header.frame_id = str(self.head_cfg.get("frame_id", "head"))
+        msg.mode = int(self.head_cfg.get("mode", 0))
+        msg.label = int(self.head_cfg.get("label", 151))
+        msg.reserved = 0
+        msg.ctrl = [motor]
+        signature = self._head_message_signature(msg)
+        window = float(self.head_cfg.get("self_command_window_s", 0.25))
+        self.expected_head_messages = [
+            expected
+            for expected in self.expected_head_messages
+            if now - expected[1] <= window
+        ][-64:]
+        self.expected_head_messages.append((signature, now))
+        self.head_publisher.publish(msg)
+        self.head_publish_count += 1
+        self.last_head_publish_at = now
+
     def _home_start_reasons(
         self,
         now: float,
@@ -2111,6 +2318,42 @@ class LocalTeleopBridge(Node):
             "last_locomotion_publish_age_s": None
             if self.last_walk_publish_at == 0.0
             else round(now - self.last_walk_publish_at, 4),
+            "head_pitch_enabled": self.head_enabled,
+            "head_pitch_binding": "left_Z + right_stick_vertical after neutral",
+            "head_pitch_left_z_pressed": self.head_z_pressed,
+            "head_pitch_active": self.head_pitch_controller.active,
+            "head_pitch_state": self.head_pitch_controller.last_transition,
+            "head_pitch_requires_neutral": (
+                self.head_pitch_controller.require_neutral
+            ),
+            "head_pitch_target_rad": self.head_pitch_controller.target_rad,
+            "head_pitch_velocity_rad_s": (
+                self.head_pitch_controller.velocity_rad_s
+            ),
+            "head_pitch_feedback_rad": self.robot_head_positions.get(
+                int(self.head_cfg.get("pitch_motor_id", 2))
+            ),
+            "head_pitch_feedback_error": self.robot_head_errors.get(
+                int(self.head_cfg.get("pitch_motor_id", 2))
+            ),
+            "head_pitch_feedback_age_s": None
+            if self.robot_head_at == 0.0
+            else round(now - self.robot_head_at, 4),
+            "head_pitch_output_reasons": self.head_output_reasons,
+            "head_pitch_limits_rad": {
+                "up": self.head_pitch_controller.min_pitch_rad,
+                "down": self.head_pitch_controller.max_pitch_rad,
+                "max_speed": self.head_pitch_controller.max_speed_rad_s,
+                "max_accel": self.head_pitch_controller.max_accel_rad_s2,
+            },
+            "head_pitch_foreign_command_age_s": None
+            if self.foreign_head_command_at == 0.0
+            else round(now - self.foreign_head_command_at, 4),
+            "head_pitch_foreign_command_count": self.foreign_head_command_count,
+            "head_pitch_publish_count": self.head_publish_count,
+            "last_head_pitch_publish_age_s": None
+            if self.last_head_publish_at == 0.0
+            else round(now - self.last_head_publish_at, 4),
             "unsupported_binding": (
                 "right_C + right_A (C+A) is not registered by xTELE 0.1.2"
             ),