meiqi/TG3
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: add session-gated TG3 data collection

Browse Source
This commit is contained in:
LengedZhao
2026-08-10 15:54:15 +08:00
parent 9d25bc9bff
commit 6caad268b8
22 changed files with 4429 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

59
tg3_local_teleop/README.md
View File

@@ -145,13 +145,70 @@ OmniSocket 进程；在此之前桥会保持最后一帧目标。未 START 或
正常待机，不触发反复重启。只有匹配的操作员 STOP 自动回 Home；意外断网不自动产生
回位运动。
## 独立数采
数采使用左摇杆按下键（L3），不是左侧 X/Y/Z 面键：
- 遥操已成功开启后，先保持 L3 松开至少 `0.5 s`；
- 连续按住 L3 `1 s` 开始一条 episode；
- 松开至少 `0.5 s`，再次连续按住 `1 s` 正常结束；
- 长按 Z+C 结束遥操、安全解除或桥退出时，也会异步请求结束当前 episode；
- 数采失败只记录错误，不解除遥操、不阻塞 50 Hz 控制，也不延迟 STOP/Home。
`button_joystick.left` 是 xTELE 5003 中独立的摇杆按压字段。已安装的 xTELE 0.1.2
没有给它注册处理函数；X/Y/Z/A/B/C 均已有厂家功能或本项目绑定，因此不复用面键。
新会话和服务重启后都先锁定为“必须松开”，畸形或陈旧按键帧不能被当作有效松开。
项目自有 `tg3-data-recorder.service` 在 Nvidia 上运行独立 `ros2 bag record`，不停止、
重配或接管 Ubuntu 厂家 `/record_bag_node`。默认仅录 `config.toml` 中的明确白名单，
包含 `/robot_state`、双臂/BrainCo 双手命令与反馈、HBWALK、IMU、电源状态以及完整
xTELE 应用帧；相机和点云默认不录，也禁止改成 `-a`。
机器人暂存目录：
```text
/home/nvidia/tg3_data_collection/
active/ # 尚未完成，不能取走
ready/ # 已正常收尾、校验并生成 READY
failed/ # 启动、磁盘、进程或校验失败，保留供诊断
```
MCAP 使用 `zstd_fast`、64 MiB cache、每 300 秒分片；每条 episode 最长 30 分钟，
启动/运行最低保留 20 GiB。正常停止以 SIGINT 让 rosbag 写完 `metadata.yaml`，随后执行
`ros2 bag info`，并确认所有 required topics 都有非零消息数；然后计算 SHA-256、写
`manifest.json` 和 `READY`，最后才原子进入 `ready/`。Nvidia 需已安装 `python3-yaml`
（现场已验证 PyYAML 6.0.1）。
部署两个机器人侧服务：
```bash
mkdir -p ~/.config/systemd/user
cp tg3-data-recorder.service tg3-local-teleop.service ~/.config/systemd/user/
systemctl --user daemon-reload
systemctl --user enable --now tg3-data-recorder.service
systemctl --user restart tg3-local-teleop.service
```
只读检查：
```bash
systemctl --user --no-pager status tg3-data-recorder.service
ros2 topic echo --once /tg3/data_collection/status
python3 -m json.tool /home/nvidia/tg3_local_teleop/status.json
find /home/nvidia/tg3_data_collection/ready -mindepth 1 -maxdepth 1 -type d
```
完成 episode 由 PS 本机服务校验后保存到项目 `Data_Get`，详见
`../tg3_data_collection/README.md`。机器人 `ready/` 中的副本不会自动删除。
## 完整重启顺序
只重启本项目的网络与机器人桥时，先启动 Nvidia 接收端，再启动 EAI 本地门控服务；
EAI 在收到物理 START 前不会建立发送 Session：
```bash
ssh nvidia@192.168.41.2 'systemctl --user restart tg3-local-teleop.service'
ssh nvidia@192.168.41.2 \
'systemctl --user restart tg3-data-recorder.service tg3-local-teleop.service'
ssh eai 'systemctl --user restart tg3-omnisocket-sender.service'
```

67
tg3_local_teleop/config.toml
View File

@@ -1,7 +1,7 @@
[network]
# Direct OmniSocket KCP input. No robot-side ZMQ receiving proxy is used.
transport = "omnisocket"
omnisocket_server = "175.178.116.187:14049"
omnisocket_server = "192.168.5.14:14049"
omnisocket_peer_id = "tg3-009027fa8190-robot"
omnisocket_expected_sender = "tg3-009027fa8190-iarm"
omnisocket_max_packet_age_ms = 300.0
@@ -55,6 +55,71 @@ joint_upper_rad = [
2.8449, 0.1920, 2.8449, 0.1920, 2.8449, 1.3265, 1.3265,
]
[data_collection]
# Left joystick press (L3) is a dedicated raw xTELE field and is not registered
# by the installed xTELE 0.1.2 button handlers. It only toggles recording while
# a validated teleoperation session is armed. Z+C STOP/safety disarm always
# ends an active recording without delaying robot control or Home.
enabled = true
button_hold_seconds = 1.0
button_release_seconds = 0.5
# This threshold qualifies only the L3 button edge; it is not a robot-motion
# watchdog and does not restore the removed 0.25 s network disarm gate.
button_input_timeout_s = 0.25
control_retry_seconds = 0.5
ack_timeout_seconds = 5.0
heartbeat_interval_seconds = 0.5
heartbeat_timeout_seconds = 3.0
status_stale_seconds = 4.0
control_topic = "/tg3/data_collection/control"
status_topic = "/tg3/data_collection/status"
iarm_frame_topic = "/tg3/data_collection/iarm_frame"
base_directory = "/home/nvidia/tg3_data_collection"
minimum_free_gib = 20.0
max_duration_seconds = 1800.0
# Explicit control-data whitelist. Camera images and point clouds are omitted
# by default to bound disk and network load; never replace this with `-a` on the
# live robot. /robot_state is the authoritative measured robot state.
topics = [
"/robot_state",
"/encoder_identical_joint",
"/arm/cmd",
"/freq_change/arm_status",
"/data_logger/arm_status",
"/left_hand/set_motor_multi",
"/right_hand/set_motor_multi",
"/left_hand/motor_status",
"/right_hand/motor_status",
"/left_hand/touch_status",
"/right_hand/touch_status",
"/hric/robot/cmd_vel",
"/hric/robot/cmd_vel_status",
"/hric/robot/rl_state",
"/imu_data",
"/power/board/key_status",
"/power/board/status",
"/power/battery/status",
"/tg3/data_collection/control",
"/tg3/data_collection/status",
"/tg3/data_collection/iarm_frame",
]
# Refuse a new episode if these core streams are absent from the live ROS
# graph. Other whitelisted streams may appear later and rosbag discovery will
# subscribe to them without restarting the recorder.
required_topics = [
"/robot_state",
"/encoder_identical_joint",
"/freq_change/arm_status",
"/left_hand/motor_status",
"/right_hand/motor_status",
"/hric/robot/rl_state",
"/tg3/data_collection/control",
"/tg3/data_collection/status",
"/tg3/data_collection/iarm_frame",
]
[locomotion]
# The two original immediate bindings are independent: right C + left-stick
# vertical controls translation; left Z + right-stick horizontal controls

1140
tg3_local_teleop/data_collection.py Executable file
View File

File diff suppressed because it is too large Load Diff

368
tg3_local_teleop/data_recorder_node.py Executable file
View File

@@ -0,0 +1,368 @@
#!/usr/bin/env python3
"""Independent ROS 2 supervisor for TG3 teleoperation data collection.
This node only controls the project-owned ``ros2 bag record`` process through
``DataRecorderManager``. It never calls, stops, or reconfigures TianGong's
factory ``record_bag_node``.
"""
from __future__ import annotations
import argparse
import json
import math
import sys
import threading
import tomllib
from dataclasses import dataclass
from pathlib import Path
from typing import Any
import rclpy
from rclpy.node import Node
from rclpy.qos import (
DurabilityPolicy,
HistoryPolicy,
QoSProfile,
ReliabilityPolicy,
)
from rclpy.utilities import remove_ros_args
from std_msgs.msg import String
from data_collection import DataRecorderManager, RecorderConfig
from data_recorder_protocol import (
RecorderControlProtocol,
topics_without_publishers,
)
DEFAULT_CONFIG = Path(__file__).with_name("config.toml")
def _boolean(value: Any, name: str) -> bool:
if type(value) is not bool:
raise ValueError(f"data_collection.{name} must be true or false")
return value
def _positive_number(value: Any, name: str) -> float:
if isinstance(value, bool) or not isinstance(value, (int, float)):
raise ValueError(f"data_collection.{name} must be a number")
result = float(value)
if not math.isfinite(result) or result <= 0.0:
raise ValueError(f"data_collection.{name} must be positive and finite")
return result
def _nonnegative_number(value: Any, name: str) -> float:
if isinstance(value, bool) or not isinstance(value, (int, float)):
raise ValueError(f"data_collection.{name} must be a number")
result = float(value)
if not math.isfinite(result) or result < 0.0:
raise ValueError(
f"data_collection.{name} must be non-negative and finite"
)
return result
def _topic(value: Any, name: str) -> str:
if (
not isinstance(value, str)
or not value.startswith("/")
or value.strip() != value
or any(character.isspace() for character in value)
):
raise ValueError(
f"data_collection.{name} must be an absolute ROS topic name"
)
return value
def _topic_list(value: Any, name: str, *, allow_empty: bool) -> tuple[str, ...]:
if not isinstance(value, list):
raise ValueError(f"data_collection.{name} must be a TOML array")
result = tuple(_topic(item, name) for item in value)
if not allow_empty and not result:
raise ValueError(f"data_collection.{name} must not be empty")
if len(set(result)) != len(result):
raise ValueError(f"data_collection.{name} contains duplicate topics")
return result
@dataclass(frozen=True)
class SupervisorSettings:
enabled: bool
control_topic: str
status_topic: str
heartbeat_timeout_seconds: float
status_publish_period_seconds: float
required_topics: tuple[str, ...]
shutdown_timeout_seconds: float
recorder: RecorderConfig
def load_settings(path: Path | str) -> SupervisorSettings:
"""Load only ``[data_collection]`` and tolerate bridge-only extra keys."""
config_path = Path(path).expanduser()
with config_path.open("rb") as stream:
document = tomllib.load(stream)
section = document.get("data_collection")
if not isinstance(section, dict):
raise ValueError("config.toml is missing [data_collection]")
enabled = _boolean(section.get("enabled", True), "enabled")
control_topic = _topic(
section.get("control_topic", "/tg3/data_collection/control"),
"control_topic",
)
status_topic = _topic(
section.get("status_topic", "/tg3/data_collection/status"),
"status_topic",
)
if control_topic == status_topic:
raise ValueError("data collection control and status topics must differ")
topics = _topic_list(section.get("topics", []), "topics", allow_empty=False)
required_topics = _topic_list(
section.get("required_topics", []),
"required_topics",
allow_empty=True,
)
unknown_required = sorted(set(required_topics) - set(topics))
if unknown_required:
raise ValueError(
"data_collection.required_topics must also appear in topics: "
+ ", ".join(unknown_required)
)
heartbeat_timeout = _positive_number(
section.get("heartbeat_timeout_seconds", 3.0),
"heartbeat_timeout_seconds",
)
status_rate = _positive_number(
section.get("status_publish_rate_hz", 5.0),
"status_publish_rate_hz",
)
minimum_free_gib = _nonnegative_number(
section.get("minimum_free_gib", 10.0), "minimum_free_gib"
)
max_duration = _positive_number(
section.get("max_duration_seconds", 1800.0),
"max_duration_seconds",
)
base_directory = section.get(
"base_directory", "/home/nvidia/tg3_data_collection"
)
if not isinstance(base_directory, str) or not base_directory.strip():
raise ValueError(
"data_collection.base_directory must be a non-empty path"
)
recorder = RecorderConfig(
base_directory=Path(base_directory).expanduser(),
topics=topics,
required_topics=required_topics,
minimum_free_bytes=int(minimum_free_gib * 1024**3),
max_duration_seconds=max_duration,
poll_interval_seconds=0.1,
sigint_timeout_seconds=15.0,
kill_timeout_seconds=3.0,
ros2_executable="ros2",
validate_bag_info=True,
bag_info_timeout_seconds=15.0,
)
return SupervisorSettings(
enabled=enabled,
control_topic=control_topic,
status_topic=status_topic,
heartbeat_timeout_seconds=heartbeat_timeout,
status_publish_period_seconds=1.0 / status_rate,
required_topics=required_topics,
# Finalization includes rosbag2 SIGINT, bag-info validation and SHA256
# generation. Leave enough time for a large, fully valid episode.
shutdown_timeout_seconds=_positive_number(
section.get("shutdown_timeout_seconds", 120.0),
"shutdown_timeout_seconds",
),
recorder=recorder,
)
class DataRecorderNode(Node):
"""Reliable String-topic adapter around ``RecorderControlProtocol``."""
def __init__(self, settings: SupervisorSettings) -> None:
super().__init__("tg3_data_recorder_supervisor")
self.settings = settings
self._closed = False
self._close_lock = threading.Lock()
self._last_logged_error = ""
recorder = DataRecorderManager(settings.recorder)
self.protocol = RecorderControlProtocol(
recorder,
heartbeat_timeout_seconds=settings.heartbeat_timeout_seconds,
start_preflight=self._start_preflight,
)
control_qos = QoSProfile(
history=HistoryPolicy.KEEP_LAST,
depth=20,
reliability=ReliabilityPolicy.RELIABLE,
durability=DurabilityPolicy.VOLATILE,
)
status_qos = QoSProfile(
history=HistoryPolicy.KEEP_LAST,
depth=10,
reliability=ReliabilityPolicy.RELIABLE,
durability=DurabilityPolicy.TRANSIENT_LOCAL,
)
self._status_publisher = self.create_publisher(
String, settings.status_topic, status_qos
)
self._control_subscription = self.create_subscription(
String,
settings.control_topic,
self._on_control,
control_qos,
)
self._timer = self.create_timer(
settings.status_publish_period_seconds, self._on_timer
)
self._publish_status(self.protocol.status())
self.get_logger().info(
"TG3 data recorder ready: control=%s status=%s base=%s topics=%d "
"heartbeat_timeout=%.3fs"
% (
settings.control_topic,
settings.status_topic,
settings.recorder.base_directory,
len(settings.recorder.topics),
settings.heartbeat_timeout_seconds,
)
)
def _start_preflight(self) -> tuple[bool, str]:
if not self.settings.required_topics:
return True, ""
try:
missing = topics_without_publishers(
self.settings.required_topics,
self.get_publishers_info_by_topic,
)
except Exception as error:
return False, f"cannot inspect ROS graph: {type(error).__name__}: {error}"
if missing:
return (
False,
"required ROS topics have no live publisher: "
+ ", ".join(missing),
)
return True, ""
def _on_control(self, message: String) -> None:
try:
status = self.protocol.handle_json(message.data)
self._publish_status(status)
except Exception as error:
# A malformed or otherwise bad data request must never kill the
# node or have any effect on the independent teleoperation bridge.
self.get_logger().error(
f"data recorder control callback failed: {type(error).__name__}: {error}"
)
def _on_timer(self) -> None:
try:
self._publish_status(self.protocol.poll())
except Exception as error:
self.get_logger().error(
f"data recorder status timer failed: {type(error).__name__}: {error}"
)
def _publish_status(self, status: dict[str, Any]) -> None:
message = String()
message.data = json.dumps(
status,
ensure_ascii=False,
sort_keys=True,
separators=(",", ":"),
allow_nan=False,
)
self._status_publisher.publish(message)
error = status.get("last_error")
error_text = error if isinstance(error, str) else ""
if error_text and error_text != self._last_logged_error:
self.get_logger().warning(f"data recorder status error: {error_text}")
self._last_logged_error = error_text
def close(self) -> bool:
with self._close_lock:
if self._closed:
return True
self._closed = True
self.get_logger().info("stopping data recorder supervisor")
try:
return self.protocol.shutdown(
timeout=self.settings.shutdown_timeout_seconds
)
except Exception as error:
self.get_logger().error(
f"data recorder shutdown failed: {type(error).__name__}: {error}"
)
return False
def _arguments(argv: list[str]) -> argparse.Namespace:
parser = argparse.ArgumentParser(
description="TG3 independent ROS 2 MCAP recorder supervisor"
)
parser.add_argument(
"--config",
type=Path,
default=DEFAULT_CONFIG,
help="TG3 bridge config.toml containing [data_collection]",
)
return parser.parse_args(remove_ros_args(argv)[1:])
def main(argv: list[str] | None = None) -> int:
process_argv = sys.argv if argv is None else argv
arguments = _arguments(process_argv)
try:
settings = load_settings(arguments.config)
except Exception as error:
print(
f"data recorder configuration error: {type(error).__name__}: {error}",
file=sys.stderr,
)
return 2
if not settings.enabled:
print("TG3 data collection is disabled in config.toml")
return 0
rclpy.init(args=process_argv)
node: DataRecorderNode | None = None
exit_code = 0
try:
node = DataRecorderNode(settings)
rclpy.spin(node)
except KeyboardInterrupt:
pass
except Exception as error:
print(
f"data recorder supervisor failed: {type(error).__name__}: {error}",
file=sys.stderr,
)
exit_code = 1
finally:
if node is not None:
if not node.close():
exit_code = 1
node.destroy_node()
if rclpy.ok():
rclpy.shutdown()
return exit_code
if __name__ == "__main__":
raise SystemExit(main())

546
tg3_local_teleop/data_recorder_protocol.py Executable file
View File

@@ -0,0 +1,546 @@
#!/usr/bin/env python3
"""Pure control protocol for the independent TG3 data recorder.
The ROS node is deliberately kept as a thin transport adapter. This module
owns JSON validation, idempotence, session/capture matching, acknowledgement
state, and the bridge-heartbeat watchdog, and can therefore be tested without
ROS installed.
"""
from __future__ import annotations
import json
import math
import time
from collections import OrderedDict
from collections.abc import Callable, Mapping, Sequence
from pathlib import Path
from typing import Any, Protocol
PROTOCOL_VERSION = 1
ACTIVE_MANAGER_STATES = {
"start_pending",
"starting",
"recording",
"stop_pending",
"finalizing",
}
class RecorderBackend(Protocol):
"""The non-blocking subset exposed by ``DataRecorderManager``."""
def request_start(
self,
session_id: str,
reason: str = "button",
*,
episode_id: str | None = None,
) -> str | None: ...
def request_stop(self, reason: str = "operator_button") -> bool: ...
def status(self) -> dict[str, Any]: ...
def shutdown(
self, *, wait: bool = True, timeout: float | None = None
) -> bool: ...
class ProtocolError(ValueError):
"""A control message is malformed or invalid for the current session."""
def __init__(self, code: str, message: str) -> None:
super().__init__(message)
self.code = code
def _nonempty_text(value: Any, name: str, *, maximum: int = 128) -> str:
if (
not isinstance(value, str)
or not value
or value.strip() != value
or len(value) > maximum
or any(ord(character) < 0x20 for character in value)
):
raise ProtocolError(
f"invalid_{name}",
f"{name} must be non-empty text up to {maximum} characters",
)
return value
def _capture_id(value: Any) -> str:
capture_id = _nonempty_text(value, "capture_id")
if (
not capture_id[0].isalnum()
or Path(capture_id).name != capture_id
or capture_id in (".", "..")
or any(
character
not in "-_.0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz"
for character in capture_id
)
):
raise ProtocolError(
"invalid_capture_id",
"capture_id must be one safe path component beginning with a letter or digit",
)
return capture_id
def _event_sequence(value: Any) -> int:
if type(value) is not int or value < 0 or value > (2**63 - 1):
raise ProtocolError(
"invalid_event_seq",
"event_seq must be an integer from 0 through 2^63-1",
)
return value
def _optional_unix_time(value: Any) -> float | None:
if value is None:
return None
if isinstance(value, bool) or not isinstance(value, (int, float)):
raise ProtocolError(
"invalid_sent_unix_s", "sent_unix_s must be a finite number"
)
result = float(value)
if not math.isfinite(result) or result < 0.0:
raise ProtocolError(
"invalid_sent_unix_s", "sent_unix_s must be a finite number"
)
return result
def parse_control_message(text: str) -> dict[str, Any]:
"""Parse one strict but forward-compatible recorder control object."""
if not isinstance(text, str) or not text or len(text.encode("utf-8")) > 16384:
raise ProtocolError(
"invalid_json", "control payload must be 1 through 16384 UTF-8 bytes"
)
try:
raw = json.loads(text)
except (json.JSONDecodeError, UnicodeError) as error:
raise ProtocolError("invalid_json", f"invalid JSON: {error}") from error
if not isinstance(raw, Mapping):
raise ProtocolError("invalid_json_object", "control JSON must be an object")
if raw.get("version") != PROTOCOL_VERSION:
raise ProtocolError(
"unsupported_version", f"version must equal {PROTOCOL_VERSION}"
)
command = raw.get("command")
if command not in ("start", "stop", "heartbeat"):
raise ProtocolError(
"invalid_command", "command must be start, stop, or heartbeat"
)
event_seq = _event_sequence(raw.get("event_seq"))
request_id = _nonempty_text(raw.get("request_id"), "request_id")
session_id = _nonempty_text(
raw.get("teleop_session_id"), "teleop_session_id"
)
capture_id = _capture_id(raw.get("capture_id"))
sent_unix_s = _optional_unix_time(raw.get("sent_unix_s"))
reason_value = raw.get("reason", "bridge_heartbeat" if command == "heartbeat" else "button")
reason = _nonempty_text(reason_value, "reason", maximum=512)
return {
"version": PROTOCOL_VERSION,
"command": command,
"event_seq": event_seq,
"request_id": request_id,
"teleop_session_id": session_id,
"capture_id": capture_id,
"reason": reason,
"sent_unix_s": sent_unix_s,
}
def topics_without_publishers(
required_topics: Sequence[str],
publisher_lookup: Callable[[str], Sequence[Any]],
) -> tuple[str, ...]:
"""Return required topics that currently have no live publisher endpoint.
ROS graph discovery can retain a topic name solely because this process is
subscribed to it. Looking up publisher endpoints is therefore the
meaningful start preflight; topic-name presence alone is insufficient.
Exceptions deliberately propagate so the caller can report graph-query
failure separately from an ordinary missing publisher.
"""
missing: list[str] = []
for topic in required_topics:
if not publisher_lookup(topic):
missing.append(topic)
return tuple(missing)
class RecorderControlProtocol:
"""Idempotent bridge-to-recorder protocol and heartbeat watchdog."""
def __init__(
self,
recorder: RecorderBackend,
*,
heartbeat_timeout_seconds: float = 3.0,
monotonic: Callable[[], float] = time.monotonic,
unix_time: Callable[[], float] = time.time,
start_preflight: Callable[[], tuple[bool, str]] | None = None,
request_cache_size: int = 512,
session_cache_size: int = 64,
) -> None:
if (
not math.isfinite(heartbeat_timeout_seconds)
or heartbeat_timeout_seconds <= 0.0
):
raise ValueError("heartbeat timeout must be positive and finite")
if request_cache_size <= 0 or session_cache_size <= 0:
raise ValueError("protocol cache sizes must be positive")
self.recorder = recorder
self.heartbeat_timeout_seconds = float(heartbeat_timeout_seconds)
self._monotonic = monotonic
self._unix_time = unix_time
self._start_preflight = start_preflight
self._request_cache_size = request_cache_size
self._session_cache_size = session_cache_size
self._request_cache: OrderedDict[str, tuple[str, bool, str, int]] = (
OrderedDict()
)
self._last_sequence_by_session: OrderedDict[str, int] = OrderedDict()
self._capture_id: str | None = None
self._session_id: str | None = None
self._last_heartbeat_monotonic: float | None = None
self._watchdog_stop_requested = False
self._last_protocol_error = ""
self._ack_event_seq: int | None = None
self._ack_request_id: str | None = None
self._ack_accepted: bool | None = None
self._ack_code = "initialized"
self._accepted_count = 0
self._rejected_count = 0
self._duplicate_count = 0
self._watchdog_stop_count = 0
def handle_json(self, text: str, *, now: float | None = None) -> dict[str, Any]:
"""Handle one payload without blocking on recorder I/O."""
timestamp = self._monotonic() if now is None else float(now)
try:
event = parse_control_message(text)
except ProtocolError as error:
self._ack_event_seq = None
self._ack_request_id = None
self._ack_accepted = False
self._ack_code = error.code
self._last_protocol_error = str(error)
self._rejected_count += 1
return self.status(now=timestamp)
fingerprint = json.dumps(event, sort_keys=True, separators=(",", ":"))
request_id = event["request_id"]
cached = self._request_cache.get(request_id)
if cached is not None:
cached_fingerprint, accepted, code, event_seq = cached
if cached_fingerprint != fingerprint:
return self._reject(
event,
"request_id_reused",
"request_id was already used for a different payload",
fingerprint=fingerprint,
cache=False,
now=timestamp,
)
self._request_cache.move_to_end(request_id)
self._duplicate_count += 1
self._set_ack(event_seq, request_id, accepted, code)
return self.status(now=timestamp)
session_id = event["teleop_session_id"]
last_sequence = self._last_sequence_by_session.get(session_id)
if last_sequence is not None and event["event_seq"] <= last_sequence:
return self._reject(
event,
"stale_event_seq",
f"event_seq must be greater than the previous value {last_sequence}",
fingerprint=fingerprint,
now=timestamp,
)
self._remember_sequence(session_id, event["event_seq"])
self._last_protocol_error = ""
try:
accepted, code = self._execute(event, timestamp)
message = (
""
if accepted
else (self._last_protocol_error or code)
)
except Exception as error:
accepted = False
code = "recorder_exception"
message = f"{type(error).__name__}: {error}"
if accepted:
self._accepted_count += 1
self._last_protocol_error = ""
else:
self._rejected_count += 1
self._last_protocol_error = message
self._remember_request(
request_id,
fingerprint,
accepted,
code,
event["event_seq"],
)
self._set_ack(
event["event_seq"],
request_id,
accepted,
code,
)
return self.status(now=timestamp)
def poll(self, *, now: float | None = None) -> dict[str, Any]:
"""Advance the watchdog and return the current status snapshot."""
timestamp = self._monotonic() if now is None else float(now)
recorder_status = self._safe_recorder_status()
manager_state = str(recorder_status.get("state", "unknown"))
if (
manager_state in ("start_pending", "starting", "recording")
and self._last_heartbeat_monotonic is not None
and not self._watchdog_stop_requested
and timestamp - self._last_heartbeat_monotonic
> self.heartbeat_timeout_seconds
):
try:
accepted = self.recorder.request_stop("bridge_heartbeat_timeout")
except Exception as error:
self._last_protocol_error = (
f"watchdog stop failed: {type(error).__name__}: {error}"
)
else:
if accepted:
self._watchdog_stop_requested = True
self._watchdog_stop_count += 1
return self._status_from_snapshot(recorder_status, timestamp)
def status(self, *, now: float | None = None) -> dict[str, Any]:
timestamp = self._monotonic() if now is None else float(now)
return self._status_from_snapshot(self._safe_recorder_status(), timestamp)
def shutdown(self, *, timeout: float | None = None) -> bool:
"""Request clean recorder finalization when the supervisor exits."""
return self.recorder.shutdown(wait=True, timeout=timeout)
def _execute(self, event: Mapping[str, Any], now: float) -> tuple[bool, str]:
command = event["command"]
if command == "start":
return self._start(event, now)
if command == "stop":
return self._stop(event)
return self._heartbeat(event, now)
def _start(self, event: Mapping[str, Any], now: float) -> tuple[bool, str]:
recorder_status = self._safe_recorder_status()
manager_state = str(recorder_status.get("state", "unknown"))
if manager_state in ACTIVE_MANAGER_STATES:
if self._matches_context(event):
return True, "already_active"
return False, "recorder_busy"
if manager_state not in ("idle",):
return False, "recorder_unavailable"
if self._start_preflight is not None:
allowed, detail = self._start_preflight()
if not allowed:
self._last_protocol_error = detail
return False, "preflight_failed"
actual_capture_id = self.recorder.request_start(
event["teleop_session_id"],
event["reason"],
episode_id=event["capture_id"],
)
if actual_capture_id is None:
return False, "recorder_busy"
if actual_capture_id != event["capture_id"]:
# A mismatched directory would make bridge samples and manifests
# impossible to correlate, so immediately fail safe by stopping.
self.recorder.request_stop("capture_id_mismatch")
return False, "capture_id_mismatch"
self._capture_id = actual_capture_id
self._session_id = event["teleop_session_id"]
self._last_heartbeat_monotonic = now
self._watchdog_stop_requested = False
return True, "start_accepted"
def _stop(self, event: Mapping[str, Any]) -> tuple[bool, str]:
recorder_status = self._safe_recorder_status()
manager_state = str(recorder_status.get("state", "unknown"))
if manager_state not in ACTIVE_MANAGER_STATES:
if self._matches_context(event) or self._capture_id is None:
return True, "already_stopped"
return False, "capture_not_active"
if not self._matches_context(event):
return False, "capture_mismatch"
accepted = self.recorder.request_stop(event["reason"])
if not accepted:
return False, "stop_rejected"
self._watchdog_stop_requested = True
return True, "stop_accepted"
def _heartbeat(
self, event: Mapping[str, Any], now: float
) -> tuple[bool, str]:
if not self._matches_context(event):
return False, "capture_mismatch"
recorder_status = self._safe_recorder_status()
manager_state = str(recorder_status.get("state", "unknown"))
if manager_state not in ACTIVE_MANAGER_STATES:
return True, "already_stopped"
self._last_heartbeat_monotonic = now
return True, "heartbeat_accepted"
def _matches_context(self, event: Mapping[str, Any]) -> bool:
return (
self._capture_id == event["capture_id"]
and self._session_id == event["teleop_session_id"]
)
def _reject(
self,
event: Mapping[str, Any],
code: str,
message: str,
*,
fingerprint: str,
cache: bool = True,
now: float,
) -> dict[str, Any]:
self._rejected_count += 1
self._last_protocol_error = message
self._set_ack(
event["event_seq"], event["request_id"], False, code
)
if cache:
self._remember_request(
event["request_id"],
fingerprint,
False,
code,
event["event_seq"],
)
return self.status(now=now)
def _remember_sequence(self, session_id: str, sequence: int) -> None:
self._last_sequence_by_session[session_id] = sequence
self._last_sequence_by_session.move_to_end(session_id)
while len(self._last_sequence_by_session) > self._session_cache_size:
self._last_sequence_by_session.popitem(last=False)
def _remember_request(
self,
request_id: str,
fingerprint: str,
accepted: bool,
code: str,
event_seq: int,
) -> None:
self._request_cache[request_id] = (
fingerprint,
accepted,
code,
event_seq,
)
self._request_cache.move_to_end(request_id)
while len(self._request_cache) > self._request_cache_size:
self._request_cache.popitem(last=False)
def _set_ack(
self,
event_seq: int,
request_id: str,
accepted: bool,
code: str,
) -> None:
self._ack_event_seq = event_seq
self._ack_request_id = request_id
self._ack_accepted = accepted
self._ack_code = code
def _safe_recorder_status(self) -> dict[str, Any]:
try:
status = self.recorder.status()
except Exception as error:
self._last_protocol_error = (
f"recorder status failed: {type(error).__name__}: {error}"
)
return {
"state": "unknown",
"recording": False,
"last_error": self._last_protocol_error,
}
return dict(status)
def _status_from_snapshot(
self, recorder_status: Mapping[str, Any], now: float
) -> dict[str, Any]:
manager_state = str(recorder_status.get("state", "unknown"))
if manager_state in ("start_pending", "starting"):
state = "starting"
elif manager_state == "recording":
state = "recording"
elif manager_state in ("stop_pending", "finalizing"):
state = "stopping"
elif manager_state == "idle":
last_result = recorder_status.get("last_result")
state = last_result if last_result in ("ready", "failed") else "idle"
else:
state = "failed"
heartbeat_age: float | None = None
if self._last_heartbeat_monotonic is not None:
heartbeat_age = max(0.0, now - self._last_heartbeat_monotonic)
manager_error = recorder_status.get("last_error")
last_error = (
str(manager_error)
if isinstance(manager_error, str) and manager_error
else self._last_protocol_error
)
return {
"version": PROTOCOL_VERSION,
"state": state,
"capture_id": self._capture_id,
"teleop_session_id": self._session_id,
"ack_event_seq": self._ack_event_seq,
"ack_request_id": self._ack_request_id,
"ack_accepted": self._ack_accepted,
"ack_code": self._ack_code,
"last_error": last_error,
"stop_reason": recorder_status.get("stop_reason"),
"heartbeat_age_s": heartbeat_age,
"heartbeat_timeout_s": self.heartbeat_timeout_seconds,
"updated_unix_s": float(self._unix_time()),
"statistics": {
"accepted": self._accepted_count,
"rejected": self._rejected_count,
"duplicates": self._duplicate_count,
"watchdog_stops": self._watchdog_stop_count,
},
"recorder": dict(recorder_status),
}
__all__ = [
"PROTOCOL_VERSION",
"ProtocolError",
"RecorderControlProtocol",
"parse_control_message",
"topics_without_publishers",
]

26
tg3_local_teleop/run_data_recorder.sh Executable file
View File

@@ -0,0 +1,26 @@
#!/usr/bin/env bash
# Load every overlay needed to deserialize the TianGong custom message types
# before rosbag2 discovers the configured topics.
set -eo pipefail
APP_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
required_setups=(
/opt/ros/jazzy/setup.bash
/home/nvidia/xos/setup.bash
/opt/robot_tele_server/install/setup.bash
"$APP_DIR/ros2_py/install/setup.bash"
)
for setup_file in "${required_setups[@]}"; do
if [[ ! -f "$setup_file" ]]; then
echo "Required ROS setup is missing: $setup_file" >&2
exit 1
fi
# shellcheck disable=SC1090
source "$setup_file"
done
set -u
exec python3 "$APP_DIR/data_recorder_node.py" \
--config "$APP_DIR/config.toml" \
"$@"

699
tg3_local_teleop/test_data_collection.py Executable file
View File

@@ -0,0 +1,699 @@
#!/usr/bin/env python3
from __future__ import annotations
import hashlib
import json
import signal
import subprocess
import tempfile
import threading
import time
import unittest
from collections.abc import Callable, Mapping, Sequence
from pathlib import Path
from typing import Any
from data_collection import (
DataRecorderManager,
RecorderConfig,
RecordingToggleGate,
left_joystick_pressed,
)
class LeftJoystickParserTest(unittest.TestCase):
def test_accepts_only_live_bool_or_binary_integer(self) -> None:
self.assertIs(
left_joystick_pressed({"button_joystick": {"left": True}}),
True,
)
self.assertIs(
left_joystick_pressed({"button_joystick": {"left": False}}),
False,
)
self.assertIs(
left_joystick_pressed({"button_joystick": {"left": 1}}), True
)
self.assertIs(
left_joystick_pressed({"button_joystick": {"left": 0}}), False
)
for value in (-1, 2, None, "true", [], 1.0, {"pressed": True}):
with self.subTest(value=value):
self.assertIsNone(
left_joystick_pressed({"button_joystick": {"left": value}})
)
def test_rejects_every_malformed_container(self) -> None:
malformed: list[Any] = [
{},
{"button_joystick": None},
{"button_joystick": []},
{"button_joystick": {}},
]
for sample in malformed:
with self.subTest(sample=sample):
self.assertIsNone(left_joystick_pressed(sample))
class RecordingToggleGateTest(unittest.TestCase):
def setUp(self) -> None:
self.gate = RecordingToggleGate(
hold_seconds=1.0, release_seconds=0.5
)
self.gate.new_session()
def stable_release(self, now: float) -> float:
self.assertIsNone(
self.gate.update(now, input_healthy=True, pressed=False)
)
now += 0.5
self.assertIsNone(
self.gate.update(now, input_healthy=True, pressed=False)
)
self.assertFalse(self.gate.require_release)
return now
def hold(self, now: float) -> tuple[float, str | None]:
self.assertIsNone(
self.gate.update(now, input_healthy=True, pressed=True)
)
now += 1.0
return now, self.gate.update(now, input_healthy=True, pressed=True)
def test_one_button_starts_and_stops_once_per_physical_hold(self) -> None:
now = self.stable_release(0.0)
now, action = self.hold(now + 0.01)
self.assertEqual(action, "start")
self.assertTrue(self.gate.active)
# Continuing the same press for any duration cannot stop recording.
self.assertIsNone(
self.gate.update(now + 20.0, input_healthy=True, pressed=True)
)
self.assertTrue(self.gate.active)
now = self.stable_release(now + 20.01)
now, action = self.hold(now + 0.01)
self.assertEqual(action, "stop")
self.assertFalse(self.gate.active)
self.assertEqual(self.gate.toggle_count, 2)
def test_new_session_is_release_locked(self) -> None:
self.assertIsNone(
self.gate.update(0.0, input_healthy=True, pressed=True)
)
self.assertIsNone(
self.gate.update(100.0, input_healthy=True, pressed=True)
)
self.assertFalse(self.gate.active)
self.assertEqual(self.gate.state, "awaiting_release")
def test_invalid_input_never_counts_as_release(self) -> None:
self.gate.update(0.0, input_healthy=True, pressed=False)
self.gate.update(0.49, input_healthy=True, pressed=False)
self.gate.update(0.5, input_healthy=True, pressed=None)
self.assertTrue(self.gate.require_release)
self.gate.update(10.0, input_healthy=True, pressed=False)
self.gate.update(10.49, input_healthy=True, pressed=False)
self.assertTrue(self.gate.require_release)
self.gate.update(10.5, input_healthy=True, pressed=False)
self.assertFalse(self.gate.require_release)
self.gate.update(10.6, input_healthy=True, pressed=True)
self.gate.update(11.59, input_healthy=False, pressed=True)
self.gate.update(20.0, input_healthy=True, pressed=True)
self.gate.update(30.0, input_healthy=True, pressed=True)
self.assertFalse(self.gate.active)
self.assertTrue(self.gate.require_release)
def test_short_press_requires_another_stable_release(self) -> None:
now = self.stable_release(0.0)
self.gate.update(now + 0.1, input_healthy=True, pressed=True)
self.gate.update(now + 0.9, input_healthy=True, pressed=False)
self.assertTrue(self.gate.require_release)
self.gate.update(now + 1.39, input_healthy=True, pressed=False)
self.assertTrue(self.gate.require_release)
self.gate.update(now + 1.4, input_healthy=True, pressed=False)
self.assertFalse(self.gate.require_release)
def test_session_end_reports_required_stop_and_relocks(self) -> None:
now = self.stable_release(0.0)
_, action = self.hold(now + 0.1)
self.assertEqual(action, "start")
self.assertEqual(self.gate.end_session(), "stop")
self.assertEqual(self.gate.state, "closed")
self.gate.new_session()
self.assertTrue(self.gate.require_release)
self.assertFalse(self.gate.active)
def test_force_inactive_handles_automatic_recorder_stop(self) -> None:
now = self.stable_release(0.0)
self.hold(now + 0.1)
self.gate.force_inactive("maximum_duration")
self.assertFalse(self.gate.active)
self.assertTrue(self.gate.require_release)
self.assertEqual(self.gate.last_transition, "maximum_duration")
class FakeProcess:
def __init__(
self,
command: Sequence[str],
*,
valid_bag: bool = True,
empty_mcap: bool = False,
unexpected_returncode: int | None = None,
ignore_sigint: bool = False,
topic_counts: Mapping[str, int] | None = None,
) -> None:
self.command = list(command)
self.signals: list[int] = []
self.returncode = unexpected_returncode
self.ignore_sigint = ignore_sigint
output = Path(self.command[self.command.index("--output") + 1])
output.mkdir(parents=True)
if valid_bag:
if topic_counts is None:
topic_counts = {"/joint_states": 25, "/tf": 50}
metadata = {
"rosbag2_bagfile_information": {
"storage_identifier": "mcap",
"topics_with_message_count": [
{
"topic_metadata": {"name": topic},
"message_count": count,
}
for topic, count in topic_counts.items()
],
}
}
(output / "metadata.yaml").write_text(
json.dumps(metadata),
encoding="utf-8",
)
(output / "data_0.mcap").write_bytes(
b"" if empty_mcap else b"fake-mcap-payload"
)
def poll(self) -> int | None:
return self.returncode
def send_signal(self, sig: int) -> None:
self.signals.append(sig)
if not self.ignore_sigint:
self.returncode = 0
def wait(self, timeout: float | None = None) -> int:
if self.returncode is None:
raise subprocess.TimeoutExpired(self.command, timeout)
return self.returncode
def terminate(self) -> None:
self.signals.append(signal.SIGTERM)
def kill(self) -> None:
self.signals.append(signal.SIGKILL)
self.returncode = -signal.SIGKILL
class FakeProcessFactory:
def __init__(self, **process_options: Any) -> None:
self.process_options = process_options
self.processes: list[FakeProcess] = []
self.commands: list[list[str]] = []
self.kwargs: list[dict[str, Any]] = []
def __call__(self, command: Sequence[str], **kwargs: Any) -> FakeProcess:
self.commands.append(list(command))
self.kwargs.append(kwargs)
process = FakeProcess(command, **self.process_options)
self.processes.append(process)
return process
class FakeBagInfoRunner:
def __init__(
self,
*,
returncode: int = 0,
stdout: str = "Files: data_0.mcap\n",
stderr: str = "",
timeout: bool = False,
) -> None:
self.returncode = returncode
self.stdout = stdout
self.stderr = stderr
self.timeout = timeout
self.calls: list[tuple[list[str], dict[str, Any]]] = []
def __call__(self, command: Sequence[str], **kwargs: Any) -> Any:
command_list = list(command)
self.calls.append((command_list, kwargs))
if self.timeout:
raise subprocess.TimeoutExpired(
command_list,
kwargs.get("timeout"),
output="partial bag info",
stderr="timed out",
)
return subprocess.CompletedProcess(
command_list,
self.returncode,
stdout=self.stdout,
stderr=self.stderr,
)
class DataRecorderManagerTest(unittest.TestCase):
def setUp(self) -> None:
self.temporary = tempfile.TemporaryDirectory()
self.base = Path(self.temporary.name) / "Data_Get"
self.managers: list[DataRecorderManager] = []
self.episode_number = 0
def tearDown(self) -> None:
for manager in self.managers:
manager.shutdown(wait=True, timeout=2.0)
self.temporary.cleanup()
def make_manager(
self,
factory: Callable[..., FakeProcess] | None = None,
*,
free_bytes: Callable[[Path], int] | None = None,
max_duration: float = 30.0,
minimum_free: int = 100,
process_options: dict[str, Any] | None = None,
command_runner: Callable[..., Any] | None = None,
validate_bag_info: bool = True,
required_topics: Sequence[str] = ("/joint_states", "/tf"),
) -> tuple[DataRecorderManager, Any]:
if factory is None:
factory = FakeProcessFactory(**(process_options or {}))
def episode_id() -> str:
self.episode_number += 1
return f"episode-{self.episode_number:03d}"
manager = DataRecorderManager(
RecorderConfig(
base_directory=self.base,
topics=("/joint_states", "/tf"),
required_topics=required_topics,
minimum_free_bytes=minimum_free,
max_duration_seconds=max_duration,
poll_interval_seconds=0.005,
sigint_timeout_seconds=0.01,
kill_timeout_seconds=0.01,
validate_bag_info=validate_bag_info,
bag_info_timeout_seconds=0.02,
),
process_factory=factory,
command_runner=command_runner or FakeBagInfoRunner(),
free_bytes=free_bytes or (lambda _path: 10_000),
episode_id_factory=episode_id,
)
self.managers.append(manager)
return manager, factory
def wait_for(
self, predicate: Callable[[], bool], timeout: float = 2.0
) -> None:
deadline = time.monotonic() + timeout
while not predicate():
if time.monotonic() >= deadline:
self.fail("condition did not become true before timeout")
time.sleep(0.002)
def test_manual_stop_uses_sigint_and_atomically_creates_ready_manifest(
self,
) -> None:
bag_info = FakeBagInfoRunner()
manager, factory = self.make_manager(command_runner=bag_info)
self.assertTrue(manager.request_start("a" * 32))
self.wait_for(lambda: manager.status()["recording"])
status = manager.status()
self.assertEqual(status["state"], "recording")
self.assertTrue(manager.request_stop("operator_button"))
self.assertTrue(manager.wait_until_idle(2.0))
status = manager.status()
self.assertEqual(status["last_result"], "ready")
self.assertEqual(status["stop_reason"], "operator_button")
ready = Path(status["last_episode_directory"])
self.assertEqual(ready.parent, self.base / "ready")
self.assertFalse((self.base / "active" / ready.name).exists())
manifest = json.loads((ready / "manifest.json").read_text())
self.assertEqual(manifest["session_id"], "a" * 32)
self.assertEqual(manifest["topics"], ["/joint_states", "/tf"])
self.assertEqual(
manifest["required_topic_message_counts"],
{"/joint_states": 25, "/tf": 50},
)
self.assertEqual(manifest["storage_id"], "mcap")
self.assertEqual(manifest["state"], "complete")
self.assertEqual(manifest["stop_reason"], "operator_button")
self.assertTrue((ready / "READY").is_file())
self.assertEqual((ready / "READY").read_text(), "ready\n")
self.assertEqual(
manifest["bag_info_validation"]["result"], "passed"
)
self.assertTrue(manifest["bag_info_validation"]["passed"])
self.assertEqual(
manifest["custom_data"],
{
"capture_id": ready.name,
"teleop_session_id": "a" * 32,
},
)
self.assertIn("started_at_utc", manifest)
self.assertIn("stopped_at_utc", manifest)
paths = {record["path"]: record for record in manifest["files"]}
self.assertEqual(
set(paths), {"bag/metadata.yaml", "bag/data_0.mcap"}
)
mcap_record = paths["bag/data_0.mcap"]
payload = (ready / "bag/data_0.mcap").read_bytes()
self.assertEqual(mcap_record["size_bytes"], len(payload))
self.assertEqual(mcap_record["size"], len(payload))
self.assertEqual(mcap_record["sha256"], hashlib.sha256(payload).hexdigest())
self.assertEqual(factory.processes[0].signals, [signal.SIGINT])
command = factory.commands[0]
self.assertEqual(
command[:5], ["ros2", "bag", "record", "--storage", "mcap"]
)
self.assertEqual(
command[command.index("--storage-preset-profile") + 1],
"zstd_fast",
)
self.assertEqual(command[command.index("--max-cache-size") + 1], "67108864")
self.assertEqual(command[command.index("--max-bag-duration") + 1], "300")
self.assertIn("--disable-keyboard-controls", command)
self.assertRegex(
command[command.index("--node-name") + 1],
r"^tg3_data_recorder_[0-9a-f]{16}$",
)
self.assertIn("--topics", command)
custom_index = command.index("--custom-data")
topics_index = command.index("--topics")
self.assertLess(custom_index, topics_index)
self.assertEqual(
command[custom_index + 1 : custom_index + 3],
[f"capture_id={ready.name}", f"teleop_session_id={'a' * 32}"],
)
self.assertEqual(command[-2:], ["/joint_states", "/tf"])
self.assertTrue(factory.kwargs[0]["start_new_session"])
self.assertEqual(
bag_info.calls[0][0],
[
"ros2",
"bag",
"info",
str(ready.parent.parent / "active" / ready.name / "bag"),
],
)
self.assertEqual(bag_info.calls[0][1]["timeout"], 0.02)
self.assertTrue(bag_info.calls[0][1]["text"])
self.assertFalse(bag_info.calls[0][1]["check"])
def test_request_methods_do_not_wait_for_slow_process_factory(self) -> None:
entered = threading.Event()
release = threading.Event()
wrapped_factory = FakeProcessFactory()
def slow_factory(command: Sequence[str], **kwargs: Any) -> FakeProcess:
entered.set()
release.wait(1.0)
return wrapped_factory(command, **kwargs)
manager, _ = self.make_manager(factory=slow_factory)
started = time.monotonic()
self.assertTrue(manager.request_start("b" * 32))
self.assertLess(time.monotonic() - started, 0.05)
self.assertTrue(entered.wait(1.0))
started = time.monotonic()
self.assertTrue(manager.request_stop("operator_button"))
self.assertLess(time.monotonic() - started, 0.05)
release.set()
self.assertTrue(manager.wait_until_idle(2.0))
def test_low_disk_space_automatically_stops_valid_episode(self) -> None:
calls = 0
def disk(_path: Path) -> int:
nonlocal calls
calls += 1
return 10_000 if calls == 1 else 0
manager, factory = self.make_manager(free_bytes=disk)
self.assertTrue(manager.request_start("c" * 32))
self.assertTrue(manager.wait_until_idle(2.0))
status = manager.status()
self.assertEqual(status["last_result"], "ready")
self.assertEqual(status["stop_reason"], "low_disk_space")
self.assertIn(signal.SIGINT, factory.processes[0].signals)
def test_maximum_duration_automatically_stops(self) -> None:
manager, _ = self.make_manager(max_duration=0.02)
self.assertTrue(manager.request_start("d" * 32))
self.assertTrue(manager.wait_until_idle(2.0))
self.assertEqual(manager.status()["last_result"], "ready")
self.assertEqual(manager.status()["stop_reason"], "maximum_duration")
def test_missing_or_empty_mcap_is_preserved_in_failed(self) -> None:
for options in ({"valid_bag": False}, {"empty_mcap": True}):
with self.subTest(options=options):
manager, _ = self.make_manager(process_options=options)
self.assertTrue(manager.request_start("e" * 32))
self.wait_for(lambda: manager.status()["recording"])
self.assertTrue(manager.request_stop())
self.assertTrue(manager.wait_until_idle(2.0))
status = manager.status()
self.assertEqual(status["last_result"], "failed")
failed = Path(status["last_episode_directory"])
self.assertEqual(failed.parent, self.base / "failed")
failure = json.loads((failed / "manifest.json").read_text())
self.assertEqual(failure["status"], "failed")
self.assertEqual(failure["session_id"], "e" * 32)
self.assertTrue(failure["error"])
manager.shutdown(timeout=2.0)
def test_unexpected_ros_exit_is_failed_not_ready(self) -> None:
manager, _ = self.make_manager(
process_options={"unexpected_returncode": 7}
)
self.assertTrue(manager.request_start("f" * 32))
self.assertTrue(manager.wait_until_idle(2.0))
status = manager.status()
self.assertEqual(status["last_result"], "failed")
self.assertIn("return code 7", status["last_error"])
self.assertEqual(
Path(status["last_episode_directory"]).parent,
self.base / "failed",
)
def test_bag_info_nonzero_exit_preserves_episode_as_failed(self) -> None:
bag_info = FakeBagInfoRunner(returncode=4, stderr="MCAP read failed")
manager, _ = self.make_manager(command_runner=bag_info)
self.assertTrue(manager.request_start("7" * 32))
self.wait_for(lambda: manager.status()["recording"])
self.assertTrue(manager.request_stop())
self.assertTrue(manager.wait_until_idle(2.0))
status = manager.status()
self.assertEqual(status["last_result"], "failed")
self.assertIn("return code 4", status["last_error"])
failed = Path(status["last_episode_directory"])
manifest = json.loads((failed / "manifest.json").read_text())
validation = manifest["bag_info_validation"]
self.assertFalse(validation["passed"])
self.assertEqual(validation["result"], "nonzero_exit")
self.assertEqual(validation["returncode"], 4)
self.assertFalse((failed / "READY").exists())
def test_missing_required_topic_is_failed_before_bag_info(self) -> None:
bag_info = FakeBagInfoRunner()
manager, _ = self.make_manager(
command_runner=bag_info,
process_options={"topic_counts": {"/joint_states": 25}},
)
self.assertTrue(manager.request_start("0" * 32))
self.wait_for(lambda: manager.status()["recording"])
self.assertTrue(manager.request_stop())
self.assertTrue(manager.wait_until_idle(2.0))
status = manager.status()
self.assertEqual(status["last_result"], "failed")
self.assertIn("missing required topics: /tf", status["last_error"])
failed = Path(status["last_episode_directory"])
manifest = json.loads((failed / "manifest.json").read_text())
self.assertEqual(
manifest["required_topic_message_counts"],
{"/joint_states": 25, "/tf": 0},
)
self.assertEqual(bag_info.calls, [])
self.assertFalse((failed / "READY").exists())
def test_zero_message_required_topic_is_failed_before_bag_info(self) -> None:
bag_info = FakeBagInfoRunner()
manager, _ = self.make_manager(
command_runner=bag_info,
process_options={
"topic_counts": {"/joint_states": 25, "/tf": 0}
},
)
self.assertTrue(manager.request_start("a0" * 16))
self.wait_for(lambda: manager.status()["recording"])
self.assertTrue(manager.request_stop())
self.assertTrue(manager.wait_until_idle(2.0))
status = manager.status()
self.assertEqual(status["last_result"], "failed")
self.assertIn("zero-message required topics: /tf", status["last_error"])
failed = Path(status["last_episode_directory"])
manifest = json.loads((failed / "manifest.json").read_text())
self.assertEqual(manifest["required_topic_message_counts"]["/tf"], 0)
self.assertEqual(bag_info.calls, [])
self.assertFalse((failed / "READY").exists())
def test_bag_info_timeout_preserves_episode_as_failed(self) -> None:
manager, _ = self.make_manager(
command_runner=FakeBagInfoRunner(timeout=True)
)
self.assertTrue(manager.request_start("8" * 32))
self.wait_for(lambda: manager.status()["recording"])
self.assertTrue(manager.request_stop())
self.assertTrue(manager.wait_until_idle(2.0))
status = manager.status()
self.assertEqual(status["last_result"], "failed")
self.assertIn("timed out", status["last_error"])
manifest = json.loads(
(
Path(status["last_episode_directory"]) / "manifest.json"
).read_text()
)
self.assertEqual(
manifest["bag_info_validation"]["result"], "timeout"
)
def test_bag_info_validation_can_be_explicitly_disabled(self) -> None:
bag_info = FakeBagInfoRunner()
manager, _ = self.make_manager(
command_runner=bag_info, validate_bag_info=False
)
self.assertTrue(manager.request_start("9" * 32))
self.wait_for(lambda: manager.status()["recording"])
self.assertTrue(manager.request_stop())
self.assertTrue(manager.wait_until_idle(2.0))
ready = Path(manager.status()["last_episode_directory"])
manifest = json.loads((ready / "manifest.json").read_text())
self.assertEqual(
manifest["bag_info_validation"]["result"], "disabled"
)
self.assertEqual(bag_info.calls, [])
def test_sigint_timeout_forces_kill_and_marks_failed(self) -> None:
manager, factory = self.make_manager(
process_options={"ignore_sigint": True}
)
self.assertTrue(manager.request_start("1" * 32))
self.wait_for(lambda: manager.status()["recording"])
self.assertTrue(manager.request_stop())
self.assertTrue(manager.wait_until_idle(2.0))
self.assertEqual(manager.status()["last_result"], "failed")
signals = factory.processes[0].signals
self.assertEqual(signals[0], signal.SIGINT)
self.assertIn(signal.SIGTERM, signals)
self.assertIn(signal.SIGKILL, signals)
def test_shutdown_waits_for_active_recording_to_finalize(self) -> None:
manager, factory = self.make_manager()
self.assertTrue(manager.request_start("2" * 32))
self.wait_for(lambda: manager.status()["recording"])
self.assertTrue(manager.shutdown(wait=True, timeout=2.0))
status = manager.status()
self.assertEqual(status["state"], "shutdown")
self.assertEqual(status["last_result"], "ready")
self.assertEqual(status["stop_reason"], "shutdown")
self.assertIn(signal.SIGINT, factory.processes[0].signals)
def test_status_snapshot_is_a_copy_and_start_is_single_flight(self) -> None:
manager, _ = self.make_manager()
first = manager.status()
first["topics"].append("/mutated")
first["state"] = "corrupt"
self.assertNotIn("/mutated", manager.status()["topics"])
self.assertEqual(manager.status()["state"], "idle")
self.assertTrue(manager.request_start("3" * 32))
self.assertFalse(manager.request_start("4" * 32))
self.wait_for(lambda: manager.status()["recording"])
self.assertTrue(manager.request_stop())
self.assertTrue(manager.wait_until_idle(2.0))
def test_caller_capture_id_is_validated_and_returned(self) -> None:
manager, _ = self.make_manager()
capture_id = "capture-20260810T120000Z-abc123"
self.assertEqual(
manager.request_start("6" * 32, episode_id=capture_id),
capture_id,
)
self.wait_for(lambda: manager.status()["recording"])
self.assertEqual(manager.status()["episode_id"], capture_id)
self.assertTrue(manager.request_stop())
self.assertTrue(manager.wait_until_idle(2.0))
self.assertEqual(
Path(manager.status()["last_episode_directory"]).name,
capture_id,
)
with self.assertRaises(ValueError):
manager.request_start("6" * 32, episode_id="../escape")
with self.assertRaises(ValueError):
manager.request_start("6" * 32, episode_id=".hidden")
with self.assertRaises(ValueError):
manager.request_start("6" * 32, episode_id="x" * 129)
with self.assertRaises(ValueError):
manager.request_start("unsafe session", episode_id="capture-safe")
def test_low_disk_before_spawn_reports_failed_without_starting_ros(self) -> None:
manager, factory = self.make_manager(free_bytes=lambda _path: 0)
self.assertTrue(manager.request_start("5" * 32))
self.assertTrue(manager.wait_until_idle(2.0))
status = manager.status()
self.assertEqual(status["last_result"], "failed")
self.assertIn("insufficient free space", status["last_error"])
self.assertEqual(factory.processes, [])
self.assertEqual(
Path(status["last_episode_directory"]).parent,
self.base / "failed",
)
failure = json.loads(
(Path(status["last_episode_directory"]) / "manifest.json").read_text()
)
self.assertEqual(failure["status"], "failed")
def test_configuration_rejects_unsafe_or_ambiguous_values(self) -> None:
with self.assertRaises(ValueError):
RecorderConfig(self.base, ())
with self.assertRaises(ValueError):
RecorderConfig(self.base, ("relative",))
with self.assertRaises(ValueError):
RecorderConfig(self.base, ("/same", "/same"))
with self.assertRaises(ValueError):
RecorderConfig(self.base, ("/ok",), minimum_free_bytes=-1)
with self.assertRaises(ValueError):
RecorderConfig(self.base, ("/ok",), validate_bag_info=1)
with self.assertRaises(ValueError):
RecorderConfig(self.base, ("/ok",), bag_info_timeout_seconds=0)
with self.assertRaises(ValueError):
RecorderConfig(
self.base,
("/recorded",),
required_topics=("/not_recorded",),
)
with self.assertRaises(ValueError):
RecorderConfig(
self.base,
("/recorded",),
required_topics=("/recorded", "/recorded"),
)
if __name__ == "__main__":
unittest.main(verbosity=2)

274
tg3_local_teleop/test_data_recorder_protocol.py Executable file
View File

@@ -0,0 +1,274 @@
#!/usr/bin/env python3
from __future__ import annotations
import json
import unittest
from typing import Any
from data_recorder_protocol import (
ProtocolError,
RecorderControlProtocol,
parse_control_message,
topics_without_publishers,
)
class FakeRecorder:
def __init__(self) -> None:
self.snapshot: dict[str, Any] = {
"state": "idle",
"recording": False,
"episode_id": None,
"session_id": None,
"stop_reason": None,
"last_result": None,
"last_error": "",
"topics": ["/robot_state"],
}
self.starts: list[tuple[str, str, str | None]] = []
self.stops: list[str] = []
self.shutdowns = 0
self.accept_start = True
def request_start(
self,
session_id: str,
reason: str = "button",
*,
episode_id: str | None = None,
) -> str | None:
self.starts.append((session_id, reason, episode_id))
if not self.accept_start or self.snapshot["state"] != "idle":
return None
self.snapshot.update(
{
"state": "start_pending",
"episode_id": episode_id,
"session_id": session_id,
"last_result": None,
}
)
return episode_id
def request_stop(self, reason: str = "operator_button") -> bool:
self.stops.append(reason)
if self.snapshot["state"] not in (
"start_pending",
"starting",
"recording",
"stop_pending",
):
return False
self.snapshot["state"] = "stop_pending"
self.snapshot["stop_reason"] = reason
return True
def status(self) -> dict[str, Any]:
return dict(self.snapshot)
def shutdown(
self, *, wait: bool = True, timeout: float | None = None
) -> bool:
self.shutdowns += 1
return True
def control(
command: str,
sequence: int,
request_id: str,
*,
capture_id: str = "capture-001",
session_id: str = "teleop-session-001",
reason: str = "test",
) -> str:
return json.dumps(
{
"version": 1,
"command": command,
"event_seq": sequence,
"request_id": request_id,
"capture_id": capture_id,
"teleop_session_id": session_id,
"reason": reason,
"sent_unix_s": 1000.0 + sequence,
}
)
class ParseControlMessageTest(unittest.TestCase):
def test_accepts_forward_compatible_object(self) -> None:
raw = json.loads(control("start", 0, "request-0"))
raw["future_field"] = {"ignored": True}
parsed = parse_control_message(json.dumps(raw))
self.assertEqual(parsed["command"], "start")
self.assertEqual(parsed["capture_id"], "capture-001")
def test_rejects_unsafe_or_ambiguous_values(self) -> None:
bad_objects = [
[],
{"version": 2},
{
"version": 1,
"command": "erase",
"event_seq": 1,
"request_id": "r",
"capture_id": "c",
"teleop_session_id": "s",
},
{
"version": 1,
"command": "start",
"event_seq": True,
"request_id": "r",
"capture_id": "../escape",
"teleop_session_id": "s",
},
]
for value in bad_objects:
with self.subTest(value=value):
with self.assertRaises(ProtocolError):
parse_control_message(json.dumps(value))
def test_required_topic_preflight_demands_a_publisher_endpoint(self) -> None:
endpoints = {
"/robot_state": [object()],
# A discovered topic with no publisher is intentionally missing.
"/encoder_identical_joint": [],
"/hric/robot/rl_state": [object(), object()],
}
calls: list[str] = []
def lookup(topic: str) -> list[object]:
calls.append(topic)
return endpoints[topic]
required = tuple(endpoints)
self.assertEqual(
topics_without_publishers(required, lookup),
("/encoder_identical_joint",),
)
self.assertEqual(calls, list(required))
class RecorderControlProtocolTest(unittest.TestCase):
def setUp(self) -> None:
self.recorder = FakeRecorder()
self.protocol = RecorderControlProtocol(
self.recorder,
heartbeat_timeout_seconds=3.0,
monotonic=lambda: 0.0,
unix_time=lambda: 1234.5,
)
def test_start_is_idempotent_and_capture_id_is_preserved(self) -> None:
message = control("start", 1, "start-request")
first = self.protocol.handle_json(message, now=10.0)
second = self.protocol.handle_json(message, now=10.2)
self.assertEqual(len(self.recorder.starts), 1)
self.assertTrue(first["ack_accepted"])
self.assertEqual(first["ack_code"], "start_accepted")
self.assertEqual(second["ack_code"], "start_accepted")
self.assertEqual(second["statistics"]["duplicates"], 1)
self.assertEqual(second["capture_id"], "capture-001")
def test_request_id_reuse_and_stale_sequences_are_rejected(self) -> None:
self.protocol.handle_json(control("start", 5, "same"), now=1.0)
reused = self.protocol.handle_json(
control("heartbeat", 6, "same"), now=1.1
)
self.assertFalse(reused["ack_accepted"])
self.assertEqual(reused["ack_code"], "request_id_reused")
stale = self.protocol.handle_json(
control("heartbeat", 4, "different"), now=1.2
)
self.assertFalse(stale["ack_accepted"])
self.assertEqual(stale["ack_code"], "stale_event_seq")
def test_matching_heartbeat_and_stop_are_immediate(self) -> None:
self.protocol.handle_json(control("start", 1, "start"), now=5.0)
self.recorder.snapshot["state"] = "recording"
heartbeat = self.protocol.handle_json(
control("heartbeat", 2, "heartbeat"), now=6.0
)
self.assertTrue(heartbeat["ack_accepted"])
self.assertEqual(heartbeat["ack_code"], "heartbeat_accepted")
stopped = self.protocol.handle_json(
control("stop", 3, "stop", reason="teleop_disarmed"), now=6.1
)
self.assertTrue(stopped["ack_accepted"])
self.assertEqual(self.recorder.stops, ["teleop_disarmed"])
def test_wrong_capture_cannot_heartbeat_or_stop_active_episode(self) -> None:
self.protocol.handle_json(control("start", 1, "start"), now=5.0)
self.recorder.snapshot["state"] = "recording"
bad_heartbeat = self.protocol.handle_json(
control(
"heartbeat",
2,
"bad-heartbeat",
capture_id="another-capture",
),
now=5.5,
)
bad_stop = self.protocol.handle_json(
control("stop", 3, "bad-stop", capture_id="another-capture"),
now=5.6,
)
self.assertFalse(bad_heartbeat["ack_accepted"])
self.assertFalse(bad_stop["ack_accepted"])
self.assertEqual(self.recorder.stops, [])
def test_missing_bridge_heartbeat_stops_recording_once(self) -> None:
self.protocol.handle_json(control("start", 1, "start"), now=10.0)
self.recorder.snapshot["state"] = "recording"
before = self.protocol.poll(now=13.0)
self.assertEqual(before["statistics"]["watchdog_stops"], 0)
after = self.protocol.poll(now=13.001)
again = self.protocol.poll(now=30.0)
self.assertEqual(self.recorder.stops, ["bridge_heartbeat_timeout"])
self.assertEqual(after["statistics"]["watchdog_stops"], 1)
self.assertEqual(again["statistics"]["watchdog_stops"], 1)
def test_manager_completion_maps_to_ready_or_failed_status(self) -> None:
self.protocol.handle_json(control("start", 1, "start"), now=1.0)
self.recorder.snapshot.update(
{
"state": "idle",
"episode_id": None,
"session_id": None,
"last_result": "ready",
}
)
ready = self.protocol.poll(now=2.0)
self.assertEqual(ready["state"], "ready")
self.assertEqual(ready["capture_id"], "capture-001")
self.recorder.snapshot.update(
{"last_result": "failed", "last_error": "bag invalid"}
)
failed = self.protocol.poll(now=2.1)
self.assertEqual(failed["state"], "failed")
self.assertEqual(failed["last_error"], "bag invalid")
def test_preflight_failure_is_status_only(self) -> None:
protocol = RecorderControlProtocol(
self.recorder,
start_preflight=lambda: (False, "missing /robot_state"),
)
status = protocol.handle_json(control("start", 1, "start"), now=1.0)
self.assertFalse(status["ack_accepted"])
self.assertEqual(status["ack_code"], "preflight_failed")
self.assertEqual(status["last_error"], "missing /robot_state")
self.assertEqual(self.recorder.starts, [])
def test_invalid_json_never_reaches_recorder(self) -> None:
status = self.protocol.handle_json("not-json", now=1.0)
self.assertFalse(status["ack_accepted"])
self.assertEqual(status["ack_code"], "invalid_json")
self.assertEqual(self.recorder.starts, [])
self.assertEqual(self.recorder.stops, [])
if __name__ == "__main__":
unittest.main()

154
tg3_local_teleop/test_session_gate.py
View File

@@ -30,6 +30,7 @@ for package, names in {
"geometry_msgs.msg": ("TwistStamped",),
"ros2_bridge_msgs.msg": ("ArmStatus",),
"sensor_msgs.msg": ("JointState",),
"std_msgs.msg": ("String",),
"std_srvs.srv": ("Trigger",),
}.items():
install_module(package, **{name: Dummy for name in names})
@@ -42,6 +43,7 @@ sys.modules[spec.name] = bridge_module
spec.loader.exec_module(bridge_module)
ArmSnapshot = bridge_module.ArmSnapshot
LocalTeleopBridge = bridge_module.LocalTeleopBridge
RecordingToggleGate = bridge_module.RecordingToggleGate
class NullLogger:
@@ -321,6 +323,158 @@ class RobotSessionGateTest(unittest.TestCase):
bridge._tick_locomotion(1.4, old_wrong_binding)
self.assertEqual(published[-1], (0.0, 0.0))
def test_l3_hold_reaches_nonblocking_capture_request(self) -> None:
bridge = LocalTeleopBridge.__new__(LocalTeleopBridge)
bridge.data_collection_enabled = True
bridge.data_collection_cfg = {
"button_input_timeout_s": 0.25,
"control_retry_seconds": 0.5,
"heartbeat_interval_seconds": 0.5,
}
bridge.data_collection_gate = RecordingToggleGate(1.0, 0.5)
bridge.data_collection_gate.new_session()
bridge.armed = True
bridge.data_pending_control = None
bridge.data_last_control_publish_at = 0.0
bridge.data_recorder_status = {}
bridge.data_recorder_status_at = 0.0
bridge.data_last_heartbeat_at = 0.0
bridge.data_last_iarm_received_at = 0.0
bridge.data_iarm_publisher = None
bridge.data_toggle_count = 0
requested: list[str] = []
bridge._request_data_capture = MethodType(
lambda _self, command, _reason: requested.append(command), bridge
)
def frame(now: float, pressed: int) -> object:
return ArmSnapshot(
{"button_joystick": {"left": pressed}}, received_at=now
)
bridge._tick_data_collection(10.0, frame(10.0, 0))
bridge._tick_data_collection(10.5, frame(10.5, 0))
bridge._tick_data_collection(10.6, frame(10.6, 1))
bridge._tick_data_collection(11.61, frame(11.61, 1))
self.assertEqual(requested, ["start"])
self.assertEqual(bridge.data_toggle_count, 1)
def test_rejected_recorder_start_clears_requested_active_state(self) -> None:
bridge = LocalTeleopBridge.__new__(LocalTeleopBridge)
bridge.data_collection_gate = RecordingToggleGate(1.0, 0.5)
bridge.data_collection_gate.new_session()
bridge.data_collection_gate.active = True
bridge.data_pending_control = {
"command": "start",
"request_id": "request-1",
"event_seq": 4,
}
bridge.data_recorder_status = {}
bridge.data_recorder_status_at = 0.0
bridge.data_capture_id = "capture-1"
bridge.data_last_transition = "start_requested"
bridge.get_logger = MethodType(lambda _self: NullLogger(), bridge)
message = Dummy()
message.data = (
'{"version":1,"state":"failed","capture_id":"capture-1",'
'"ack_request_id":"request-1","ack_event_seq":4,'
'"ack_accepted":false,"ack_code":"preflight_failed",'
'"last_error":"required topic missing"}'
)
bridge._on_data_recorder_status(message)
self.assertIsNone(bridge.data_pending_control)
self.assertFalse(bridge.data_collection_gate.active)
self.assertIn("start_rejected", bridge.data_last_transition)
def test_supervisor_context_loss_clears_false_recording_state(self) -> None:
bridge = LocalTeleopBridge.__new__(LocalTeleopBridge)
bridge.data_collection_gate = RecordingToggleGate(1.0, 0.5)
bridge.data_collection_gate.new_session()
bridge.data_collection_gate.active = True
bridge.data_pending_control = None
bridge.data_recorder_status = {}
bridge.data_recorder_status_at = 0.0
bridge.data_capture_id = "capture-old"
bridge.data_last_transition = "recording"
bridge.get_logger = MethodType(lambda _self: NullLogger(), bridge)
message = Dummy()
message.data = (
'{"version":1,"state":"idle","capture_id":null,'
'"ack_request_id":null,"ack_event_seq":null,'
'"ack_accepted":null,"last_error":""}'
)
bridge._on_data_recorder_status(message)
self.assertFalse(bridge.data_collection_gate.active)
self.assertEqual(
bridge.data_last_transition,
"recorder context was lost or replaced",
)
def test_missing_supervisor_ack_cancels_false_start(self) -> None:
bridge = LocalTeleopBridge.__new__(LocalTeleopBridge)
bridge.data_collection_enabled = True
bridge.data_collection_cfg = {
"control_retry_seconds": 0.5,
"ack_timeout_seconds": 5.0,
"heartbeat_interval_seconds": 0.5,
"status_stale_seconds": 4.0,
}
bridge.data_collection_gate = RecordingToggleGate(1.0, 0.5)
bridge.data_collection_gate.new_session()
bridge.data_collection_gate.active = True
bridge.armed = False
bridge.data_pending_control = {"command": "start"}
bridge.data_pending_control_since = 10.0
bridge.data_last_control_publish_at = 10.0
bridge.data_control_publisher = None
bridge.data_recorder_status = {}
bridge.data_recorder_status_at = 0.0
bridge.data_last_heartbeat_at = 0.0
bridge.data_last_iarm_received_at = 0.0
bridge.data_iarm_publisher = None
bridge.data_last_transition = "start_requested"
bridge.get_logger = MethodType(lambda _self: NullLogger(), bridge)
bridge._tick_data_collection(15.0, None)
self.assertIsNone(bridge.data_pending_control)
self.assertFalse(bridge.data_collection_gate.active)
self.assertEqual(bridge.data_last_transition, "start_ack_timeout")
def test_stale_recorder_status_requests_capture_stop(self) -> None:
bridge = LocalTeleopBridge.__new__(LocalTeleopBridge)
bridge.data_collection_enabled = True
bridge.data_collection_cfg = {
"control_retry_seconds": 0.5,
"ack_timeout_seconds": 5.0,
"heartbeat_interval_seconds": 0.5,
"status_stale_seconds": 4.0,
}
bridge.data_collection_gate = RecordingToggleGate(1.0, 0.5)
bridge.data_collection_gate.new_session()
bridge.data_collection_gate.active = True
bridge.armed = False
bridge.data_pending_control = None
bridge.data_pending_control_since = 0.0
bridge.data_last_control_publish_at = 0.0
bridge.data_recorder_status = {"state": "recording"}
bridge.data_recorder_status_at = 10.0
bridge.data_last_heartbeat_at = 10.0
bridge.data_last_iarm_received_at = 0.0
bridge.data_iarm_publisher = None
bridge.data_last_transition = "recording"
bridge.get_logger = MethodType(lambda _self: NullLogger(), bridge)
requested: list[str] = []
bridge._request_data_capture = MethodType(
lambda _self, command, _reason: requested.append(command), bridge
)
bridge._tick_data_collection(14.0, None)
self.assertFalse(bridge.data_collection_gate.active)
self.assertEqual(requested, ["stop"])
self.assertEqual(bridge.data_last_transition, "recorder_status_stale")
if __name__ == "__main__":
unittest.main()

23
tg3_local_teleop/tg3-data-recorder.service Normal file
View File

@@ -0,0 +1,23 @@
[Unit]
Description=TG3 independent teleoperation MCAP recorder supervisor
After=network-online.target
Wants=network-online.target
Before=tg3-local-teleop.service
[Service]
Type=simple
WorkingDirectory=/home/nvidia/tg3_local_teleop
UMask=0027
ExecStartPre=/home/nvidia/tg3_local_teleop/wait_ros_ready.sh
ExecStart=/home/nvidia/tg3_local_teleop/run_data_recorder.sh
Restart=on-failure
RestartSec=2
KillSignal=SIGINT
# Signal only the Python supervisor first; it owns clean SIGINT/finalization
# of rosbag2. The whole cgroup is still killed if the generous deadline ends.
KillMode=mixed
TimeoutStartSec=150
TimeoutStopSec=150
[Install]
WantedBy=default.target

4
tg3_local_teleop/tg3-local-teleop.service
View File

@@ -1,7 +1,7 @@
[Unit]
Description=TG3 local TS1P dual-arm bridge (no cloud pairing)
After=network-online.target
Wants=network-online.target
After=network-online.target tg3-data-recorder.service
Wants=network-online.target tg3-data-recorder.service
[Service]
Type=simple

345
tg3_local_teleop/tg3_local_teleop.py
View File

@@ -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,