feat: add session-gated TG3 data collection

tg3_local_teleop/tg3_local_teleop.py

@@ -16,6 +16,7 @@ import struct
 import threading
 import time
 import tomllib
+import uuid
 from dataclasses import dataclass
 from pathlib import Path
 from typing import Any
@@ -28,8 +29,10 @@ from geometry_msgs.msg import TwistStamped
 from rclpy.node import Node
 from ros2_bridge_msgs.msg import ArmStatus
 from sensor_msgs.msg import JointState
+from std_msgs.msg import String
 from std_srvs.srv import Trigger
 
+from data_collection import RecordingToggleGate, left_joystick_pressed
 from gesture_toggle import (
     GestureToggle,
     normalized_pose_to_positions,
@@ -426,6 +429,32 @@ class LocalTeleopBridge(Node):
         self.locomotion_enabled = bool(
             self.locomotion_cfg.get("enabled", False)
         )
+        self.data_collection_cfg = config.get("data_collection", {})
+        self.data_collection_enabled = bool(
+            self.data_collection_cfg.get("enabled", False)
+        )
+        self.data_collection_gate = RecordingToggleGate(
+            hold_seconds=float(
+                self.data_collection_cfg.get("button_hold_seconds", 1.0)
+            ),
+            release_seconds=float(
+                self.data_collection_cfg.get("button_release_seconds", 0.5)
+            ),
+        )
+        self.data_control_publisher = None
+        self.data_iarm_publisher = None
+        self.data_recorder_status: dict[str, Any] = {}
+        self.data_recorder_status_at = 0.0
+        self.data_pending_control: dict[str, Any] | None = None
+        self.data_pending_control_since = 0.0
+        self.data_last_control_publish_at = 0.0
+        self.data_last_heartbeat_at = 0.0
+        self.data_last_iarm_received_at = 0.0
+        self.data_capture_id: str | None = None
+        self.data_capture_session_id: str | None = None
+        self.data_event_seq = 0
+        self.data_toggle_count = 0
+        self.data_last_transition = "initialized; waiting for an armed session"
         self.hand_publishers: dict[str, Any] = {}
         self.robot_hand_positions: dict[str, list[int] | None] = {
             side: None for side in HAND_SIDES
@@ -492,6 +521,23 @@ class LocalTeleopBridge(Node):
         self.create_service(
             Trigger, ros_cfg["cancel_home_service"], self._on_cancel_home_request
         )
+        if getattr(self, "data_collection_enabled", False):
+            self.data_control_publisher = self.create_publisher(
+                String,
+                str(self.data_collection_cfg["control_topic"]),
+                10,
+            )
+            self.data_iarm_publisher = self.create_publisher(
+                String,
+                str(self.data_collection_cfg["iarm_frame_topic"]),
+                10,
+            )
+            self.create_subscription(
+                String,
+                str(self.data_collection_cfg["status_topic"]),
+                self._on_data_recorder_status,
+                10,
+            )
 
         self.source = LatestArmData(net_cfg)
         self.source.start()
@@ -526,7 +572,8 @@ class LocalTeleopBridge(Node):
         self.get_logger().info(
             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"locomotion={self.locomotion_enabled}; "
+            f"data_collection={self.data_collection_enabled}"
         )
 
     def close(self) -> None:
@@ -558,6 +605,242 @@ class LocalTeleopBridge(Node):
         self.robot_hand_states[side] = states
         self.robot_hand_at[side] = time.monotonic()
 
+    def _on_data_recorder_status(self, msg: String) -> None:
+        """Consume recorder acknowledgements without affecting robot control."""
+
+        try:
+            status = json.loads(msg.data)
+            if not isinstance(status, dict) or status.get("version") != 1:
+                raise ValueError("unsupported recorder status")
+        except (TypeError, ValueError, json.JSONDecodeError) as exc:
+            self.get_logger().warning(f"ignored invalid data-recorder status: {exc}")
+            return
+
+        self.data_recorder_status = status
+        self.data_recorder_status_at = time.monotonic()
+        pending = self.data_pending_control
+        if (
+            pending is not None
+            and status.get("ack_request_id") == pending.get("request_id")
+            and status.get("ack_event_seq") == pending.get("event_seq")
+        ):
+            accepted = status.get("ack_accepted")
+            if type(accepted) is bool:
+                self.data_pending_control = None
+                self.data_pending_control_since = 0.0
+                if not accepted:
+                    detail = str(
+                        status.get("last_error")
+                        or status.get("ack_code")
+                        or "recorder rejected request"
+                    )
+                    if pending.get("command") == "start":
+                        self.data_collection_gate.force_inactive(
+                            f"recorder_rejected: {detail}"
+                        )
+                    self.data_last_transition = (
+                        f"{pending.get('command')}_rejected: {detail}"
+                    )
+                    self.get_logger().error(
+                        "data recorder rejected %s: %s"
+                        % (pending.get("command"), detail)
+                    )
+
+        state = status.get("state")
+        if (
+            state in ("ready", "failed")
+            and status.get("capture_id") == self.data_capture_id
+            and self.data_collection_gate.active
+        ):
+            self.data_collection_gate.force_inactive(
+                f"recorder_{state}: {status.get('last_error', '')}".rstrip()
+            )
+            self.data_last_transition = self.data_collection_gate.last_transition
+        elif (
+            self.data_pending_control is None
+            and self.data_collection_gate.active
+            and state in ("idle", "ready", "failed", "recording", "stopping")
+            and status.get("capture_id") != self.data_capture_id
+        ):
+            # A restarted supervisor has no in-memory context for the old bag.
+            # Reconcile the UI gate instead of displaying a false recording
+            # state forever; the supervisor owns cleanup of its old cgroup and
+            # active/failed directory.
+            self.data_collection_gate.force_inactive(
+                "recorder context was lost or replaced"
+            )
+            self.data_last_transition = self.data_collection_gate.last_transition
+            self.get_logger().error(
+                "data recorder no longer owns the requested capture; "
+                "release L3 before starting a new episode"
+            )
+
+    def _publish_data_control(self, payload: dict[str, Any]) -> None:
+        if self.data_control_publisher is None:
+            return
+        try:
+            message = String()
+            message.data = json.dumps(
+                payload, ensure_ascii=False, separators=(",", ":")
+            )
+            self.data_control_publisher.publish(message)
+            self.data_last_control_publish_at = time.monotonic()
+        except Exception as exc:  # Data collection must never stop robot control.
+            self.get_logger().error(f"cannot publish data-recorder control: {exc}")
+
+    def _request_data_capture(self, command: str, reason: str) -> None:
+        if not self.data_collection_enabled:
+            return
+        if command == "start":
+            self.data_capture_id = (
+                time.strftime("%Y%m%dT%H%M%SZ", time.gmtime())
+                + "_TG3_"
+                + uuid.uuid4().hex[:8]
+            )
+        if self.data_capture_id is None:
+            return
+        if self.data_capture_session_id is None:
+            self.data_capture_session_id = self.active_session_id or (
+                "direct_" + uuid.uuid4().hex
+            )
+
+        self.data_event_seq += 1
+        payload: dict[str, Any] = {
+            "version": 1,
+            "command": command,
+            "event_seq": self.data_event_seq,
+            "request_id": uuid.uuid4().hex,
+            "capture_id": self.data_capture_id,
+            "teleop_session_id": self.data_capture_session_id,
+            "reason": reason,
+            "sent_unix_s": time.time(),
+        }
+        self.data_pending_control = payload
+        self.data_pending_control_since = time.monotonic()
+        self._publish_data_control(payload)
+        self.data_last_transition = f"{command}_requested: {reason}"
+
+    def _publish_data_heartbeat(self) -> None:
+        if self.data_capture_id is None or self.data_capture_session_id is None:
+            return
+        self.data_event_seq += 1
+        payload = {
+            "version": 1,
+            "command": "heartbeat",
+            "event_seq": self.data_event_seq,
+            "request_id": uuid.uuid4().hex,
+            "capture_id": self.data_capture_id,
+            "teleop_session_id": self.data_capture_session_id,
+            "sent_unix_s": time.time(),
+        }
+        self._publish_data_control(payload)
+        self.data_last_heartbeat_at = time.monotonic()
+
+    def _tick_data_collection(
+        self, now: float, sample: ArmSnapshot | None
+    ) -> None:
+        if not self.data_collection_enabled:
+            return
+
+        if self.armed:
+            input_timeout = float(
+                self.data_collection_cfg.get("button_input_timeout_s", 0.25)
+            )
+            input_healthy = (
+                sample is not None and now - sample.received_at <= input_timeout
+            )
+            pressed = (
+                None if sample is None else left_joystick_pressed(sample.data)
+            )
+            action = self.data_collection_gate.update(
+                now,
+                input_healthy=input_healthy,
+                pressed=pressed,
+            )
+            if action is not None:
+                self.data_toggle_count += 1
+                self._request_data_capture(
+                    action,
+                    "left joystick press held for "
+                    f"{self.data_collection_gate.hold_seconds:.1f}s",
+                )
+
+        retry_s = float(self.data_collection_cfg.get("control_retry_seconds", 0.5))
+        if (
+            self.data_pending_control is not None
+            and now - self.data_last_control_publish_at >= retry_s
+        ):
+            self._publish_data_control(self.data_pending_control)
+
+        ack_timeout_s = float(
+            self.data_collection_cfg.get("ack_timeout_seconds", 5.0)
+        )
+        if (
+            self.data_pending_control is not None
+            and self.data_pending_control_since > 0.0
+            and now - self.data_pending_control_since >= ack_timeout_s
+        ):
+            expired_command = str(self.data_pending_control.get("command"))
+            self.data_pending_control = None
+            self.data_pending_control_since = 0.0
+            if expired_command == "start":
+                self.data_collection_gate.force_inactive(
+                    "recorder START acknowledgement timed out"
+                )
+            self.data_last_transition = f"{expired_command}_ack_timeout"
+            self.get_logger().error(
+                f"data recorder {expired_command} acknowledgement timed out"
+            )
+
+        status_stale_s = float(
+            self.data_collection_cfg.get("status_stale_seconds", 4.0)
+        )
+        if (
+            self.data_collection_gate.active
+            and self.data_pending_control is None
+            and self.data_recorder_status_at > 0.0
+            and now - self.data_recorder_status_at >= status_stale_s
+        ):
+            self.data_collection_gate.force_inactive(
+                "recorder status heartbeat became stale"
+            )
+            self.data_last_transition = "recorder_status_stale"
+            self._request_data_capture(
+                "stop", "recorder status heartbeat became stale"
+            )
+            self.get_logger().error(
+                "data recorder status became stale; capture stop requested"
+            )
+
+        recorder_state = self.data_recorder_status.get("state")
+        heartbeat_s = float(
+            self.data_collection_cfg.get("heartbeat_interval_seconds", 0.5)
+        )
+        if (
+            self.data_collection_gate.active
+            and self.data_pending_control is None
+            and recorder_state in ("starting", "recording")
+            and now - self.data_last_heartbeat_at >= heartbeat_s
+        ):
+            self._publish_data_heartbeat()
+
+        if (
+            self.data_collection_gate.active
+            and sample is not None
+            and sample.received_at != self.data_last_iarm_received_at
+            and self.data_iarm_publisher is not None
+        ):
+            try:
+                message = String()
+                # Preserve the complete xTELE/OmniSocket application frame.
+                message.data = json.dumps(
+                    sample.data, ensure_ascii=False, separators=(",", ":")
+                )
+                self.data_iarm_publisher.publish(message)
+                self.data_last_iarm_received_at = sample.received_at
+            except Exception as exc:
+                self.get_logger().error(f"cannot publish xTELE capture frame: {exc}")
+
     @staticmethod
     def _hand_message_signature(msg: SetMotorMulti) -> tuple[Any, ...]:
         return (
@@ -690,6 +973,11 @@ class LocalTeleopBridge(Node):
                 success=False,
             )
 
+        # The recorder toggle is deliberately downstream of every disarm
+        # 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.right_point_gesture_enabled:
             gesture_toggled = self.right_point_gesture.update(
                 now,
@@ -937,6 +1225,10 @@ class LocalTeleopBridge(Node):
         self.active_session_id = session_id
         if self.right_point_gesture_enabled:
             self.right_point_gesture.new_session()
+        if getattr(self, "data_collection_enabled", False):
+            self.data_capture_session_id = session_id
+            self.data_collection_gate.new_session()
+            self.data_last_transition = "new teleoperation session"
         self.armed = True
         # Start both slew limiters at measured robot feedback, never at a
         # potentially distant first network target.
@@ -1002,6 +1294,10 @@ class LocalTeleopBridge(Node):
             return
         if self.right_point_gesture_enabled:
             self.right_point_gesture.new_session()
+        if getattr(self, "data_collection_enabled", False):
+            self.data_capture_session_id = "direct_" + uuid.uuid4().hex
+            self.data_collection_gate.new_session()
+            self.data_last_transition = "new direct-LAN teleoperation session"
         self.armed = True
         # Start the slew limiter at measured robot feedback.  Using None here
         # would make the first armed frame jump directly to the TS1P target.
@@ -1020,6 +1316,13 @@ class LocalTeleopBridge(Node):
 
     def _disarm(self, reason: str) -> None:
         was_armed = self.armed
+        if getattr(self, "data_collection_enabled", False):
+            action = self.data_collection_gate.end_session(reason)
+            self.data_last_transition = self.data_collection_gate.last_transition
+            if action == "stop":
+                # This is a non-blocking ROS request.  It never delays robot
+                # disarm, locomotion zeroing, or the existing Home sequence.
+                self._request_data_capture("stop", reason)
         self._stop_locomotion(reason)
         if self.right_point_gesture_enabled:
             # Clearing the logical override must not publish a hand target.
@@ -1550,6 +1853,46 @@ class LocalTeleopBridge(Node):
             "custom_hand_feedback_policy": (
                 "startup_qualifies_teleop; runtime_gap_pauses_hands_only"
             ),
+            "data_collection_enabled": self.data_collection_enabled,
+            "data_collection_binding": (
+                "left joystick press (L3), hold "
+                f"{self.data_collection_gate.hold_seconds:.1f}s toggle; "
+                f"release {self.data_collection_gate.release_seconds:.1f}s"
+            ),
+            "data_collection_gate_state": self.data_collection_gate.state,
+            "data_collection_requested_active": self.data_collection_gate.active,
+            "data_collection_requires_release": (
+                self.data_collection_gate.require_release
+            ),
+            "data_collection_button_pressed": (
+                self.data_collection_gate.button_pressed
+            ),
+            "data_collection_hold_s": (
+                0.0
+                if self.data_collection_gate.hold_started_at is None
+                else round(now - self.data_collection_gate.hold_started_at, 2)
+            ),
+            "data_collection_capture_id": self.data_capture_id,
+            "data_collection_session_id": self.data_capture_session_id,
+            "data_collection_toggle_count": self.data_toggle_count,
+            "data_collection_last_transition": self.data_last_transition,
+            "data_collection_pending_command": (
+                None
+                if self.data_pending_control is None
+                else self.data_pending_control.get("command")
+            ),
+            "data_collection_pending_age_s": (
+                None
+                if self.data_pending_control is None
+                or self.data_pending_control_since == 0.0
+                else round(now - self.data_pending_control_since, 3)
+            ),
+            "data_recorder_status": self.data_recorder_status,
+            "data_recorder_status_age_s": (
+                None
+                if self.data_recorder_status_at == 0.0
+                else round(now - self.data_recorder_status_at, 3)
+            ),
             "runtime_hand_output_ready": self.hand_output_ready,
             "runtime_hand_output_reasons": self.runtime_hand_output_reasons,
             "safety_ready": not reasons,