#!/usr/bin/env python3 """Decode TG3 joint streams from one locally stored MCAP episode into CSV.""" from __future__ import annotations import argparse import csv import json from pathlib import Path from typing import Any, Callable ARM_MOTOR_IDS = [*range(11, 18), *range(21, 28)] JOINT_NAMES = [ *(f"left_joints_{index}" for index in range(7)), *(f"right_joints_{index}" for index in range(7)), ] CUSTOM_TOPICS = { "/robot_state": "ros2_bridge_msgs/msg/RobotState", "/freq_change/arm_status": "ros2_bridge_msgs/msg/ArmStatus", "/data_logger/arm_status": "ros2_bridge_msgs/msg/ArmStatus", "/arm/cmd": "ros2_bridge_msgs/msg/ArmCtrl", } TOPICS = [ "/robot_state", "/encoder_identical_joint", "/freq_change/arm_status", "/data_logger/arm_status", "/arm/cmd", "/tg3/data_collection/iarm_frame", ] def _message_directory() -> Path: return ( Path(__file__).resolve().parents[1] / "tg3_local_teleop" / "ros2_py" / "src" / "ros2_bridge_msgs" / "msg" ) def _read_definition(name: str) -> str: path = _message_directory() / f"{name}.msg" if not path.is_file(): raise RuntimeError(f"required message definition is missing: {path}") return path.read_text(encoding="utf-8") def _definition_bundle(root_name: str, dependencies: list[tuple[str, str]]) -> str: if root_name == "ArmCtrl": text = ( "std_msgs/Header header\n" "uint8 mode\n" "uint8 label\n" "uint8 reserved\n" "ros2_bridge_msgs/MotorCtrl[] ctrl\n" ) else: text = _read_definition(root_name) for full_name, definition in dependencies: text += f"\n===\nMSG: {full_name}\n{definition}" return text def _custom_decoders() -> dict[str, Callable[[bytes], Any]]: try: from mcap_ros2._dynamic import generate_dynamic except ImportError as exc: raise RuntimeError( "missing MCAP decoder; install with: pip install mcap mcap-ros2-support" ) from exc header = "builtin_interfaces/Time stamp\nstring frame_id\n" motor_status = _read_definition("MotorStatus") common = [ ("std_msgs/msg/Header", header), ("ros2_bridge_msgs/msg/MotorStatus", motor_status), ] robot_dependencies = common + [ (f"ros2_bridge_msgs/msg/{name}", _read_definition(name)) for name in ( "HeadStatus", "WaistStatus", "LegStatus", "ArmStatus", "ImuStatus", ) ] schemas = { "ros2_bridge_msgs/msg/RobotState": _definition_bundle( "RobotState", robot_dependencies ), "ros2_bridge_msgs/msg/ArmStatus": _definition_bundle( "ArmStatus", common ), "ros2_bridge_msgs/msg/ArmCtrl": _definition_bundle( "ArmCtrl", [ ("std_msgs/msg/Header", header), ("ros2_bridge_msgs/msg/MotorCtrl", _read_definition("MotorCtrl")), ], ), } return { type_name: generate_dynamic(type_name, text)[type_name] for type_name, text in schemas.items() } def _stamp_ns(header: Any) -> int: return int(header.stamp.sec) * 1_000_000_000 + int(header.stamp.nanosec) def _side_and_index(motor_id: int) -> tuple[str, int]: if 11 <= motor_id <= 17: return "left", motor_id - 11 if 21 <= motor_id <= 27: return "right", motor_id - 21 return "unknown", -1 def _open_csv(path: Path, columns: list[str]) -> tuple[Any, csv.DictWriter]: stream = path.open("x", encoding="utf-8", newline="") writer = csv.DictWriter(stream, fieldnames=columns) writer.writeheader() return stream, writer def decode_episode(episode: Path, output: Path) -> dict[str, int]: try: from mcap.reader import make_reader from mcap_ros2.decoder import DecoderFactory except ImportError as exc: raise RuntimeError( "missing MCAP decoder; install with: pip install mcap mcap-ros2-support" ) from exc bag_files = sorted((episode / "bag").glob("*.mcap")) if not bag_files: raise RuntimeError(f"no MCAP files found under {episode / 'bag'}") output.mkdir(parents=True, exist_ok=False) feedback_columns = [ "bag_time_ns", "ros_time_ns", "topic", "side", "joint_index", "motor_id", "position_rad", "speed_rad_s", "current_a", "temperature_c", "mos_temperature_c", "error", ] target_columns = [ "bag_time_ns", "ros_time_ns", "joint_index", "joint_name", "motor_id", "side", "position_rad", "velocity_rad_s", "effort", ] command_columns = [ "bag_time_ns", "ros_time_ns", "mode", "label", "motor_id", "side", "joint_index", "kp", "kd", "position_rad", "speed_rad_s", "torque_feedforward", "current_limit_a", ] source_columns = [ "bag_time_ns", "source_time_ms", "joint_index", "joint_name", "motor_id", "side", "position_rad", ] streams: list[Any] = [] feedback_stream, feedback_writer = _open_csv( output / "robot_arm_feedback.csv", feedback_columns ) target_stream, target_writer = _open_csv( output / "teleop_joint_target.csv", target_columns ) command_stream, command_writer = _open_csv( output / "vendor_arm_command.csv", command_columns ) source_stream, source_writer = _open_csv( output / "iarm_source_joint.csv", source_columns ) streams.extend((feedback_stream, target_stream, command_stream, source_stream)) counts = { "robot_arm_feedback_rows": 0, "teleop_joint_target_rows": 0, "vendor_arm_command_rows": 0, "iarm_source_joint_rows": 0, } custom_decoders = _custom_decoders() standard_factory = DecoderFactory() try: for bag_file in bag_files: with bag_file.open("rb") as bag_stream: reader = make_reader(bag_stream) for schema, channel, message in reader.iter_messages(topics=TOPICS): topic = channel.topic if topic in CUSTOM_TOPICS: decoded = custom_decoders[CUSTOM_TOPICS[topic]](message.data) else: decoder = standard_factory.decoder_for( channel.message_encoding, schema ) if decoder is None: raise RuntimeError( f"no decoder for {topic} ({schema.name if schema else 'no schema'})" ) decoded = decoder(message.data) if topic in ( "/robot_state", "/freq_change/arm_status", "/data_logger/arm_status", ): status_message = decoded.arm if topic == "/robot_state" else decoded for motor in status_message.status: motor_id = int(motor.name) side, joint_index = _side_and_index(motor_id) feedback_writer.writerow( { "bag_time_ns": message.log_time, "ros_time_ns": _stamp_ns(decoded.header), "topic": topic, "side": side, "joint_index": joint_index, "motor_id": motor_id, "position_rad": motor.pos, "speed_rad_s": motor.speed, "current_a": motor.current, "temperature_c": motor.temperature, "mos_temperature_c": motor.mos_temperature, "error": motor.error, } ) counts["robot_arm_feedback_rows"] += 1 elif topic == "/encoder_identical_joint": for index, position in enumerate(decoded.position): motor_id = ARM_MOTOR_IDS[index] if index < 14 else -1 side, _ = _side_and_index(motor_id) target_writer.writerow( { "bag_time_ns": message.log_time, "ros_time_ns": _stamp_ns(decoded.header), "joint_index": index, "joint_name": decoded.name[index] if index < len(decoded.name) else "", "motor_id": motor_id, "side": side, "position_rad": position, "velocity_rad_s": decoded.velocity[index] if index < len(decoded.velocity) else "", "effort": decoded.effort[index] if index < len(decoded.effort) else "", } ) counts["teleop_joint_target_rows"] += 1 elif topic == "/arm/cmd": for motor in decoded.ctrl: motor_id = int(motor.name) side, joint_index = _side_and_index(motor_id) command_writer.writerow( { "bag_time_ns": message.log_time, "ros_time_ns": _stamp_ns(decoded.header), "mode": decoded.mode, "label": decoded.label, "motor_id": motor_id, "side": side, "joint_index": joint_index, "kp": motor.kp, "kd": motor.kd, "position_rad": motor.pos, "speed_rad_s": motor.spd, "torque_feedforward": motor.tor, "current_limit_a": motor.cur, } ) counts["vendor_arm_command_rows"] += 1 elif topic == "/tg3/data_collection/iarm_frame": payload = json.loads(decoded.data) positions = payload.get("arm", {}).get("position", {}) values = [ *positions.get("left", []), *positions.get("right", []), ] if len(values) != 14: continue source_time = payload.get("timestamp", "") for index, position in enumerate(values): motor_id = ARM_MOTOR_IDS[index] side, _ = _side_and_index(motor_id) source_writer.writerow( { "bag_time_ns": message.log_time, "source_time_ms": source_time, "joint_index": index, "joint_name": JOINT_NAMES[index], "motor_id": motor_id, "side": side, "position_rad": position, } ) counts["iarm_source_joint_rows"] += 1 finally: for stream in streams: stream.close() (output / "summary.json").write_text( json.dumps( { "episode": episode.name, "joint_order": JOINT_NAMES, "motor_ids": ARM_MOTOR_IDS, "units": { "position": "rad", "speed": "rad/s", "current": "A", "temperature": "degC", }, "rows": counts, }, ensure_ascii=False, indent=2, ) + "\n", encoding="utf-8", ) return counts def main() -> int: parser = argparse.ArgumentParser() parser.add_argument("episode", type=Path, help="Data_Get episode directory") parser.add_argument( "--output", type=Path, help="new output directory (default: /decoded_joints)", ) args = parser.parse_args() episode = args.episode.resolve() output = (args.output or episode / "decoded_joints").resolve() counts = decode_episode(episode, output) print(json.dumps({"output": str(output), "rows": counts}, indent=2)) return 0 if __name__ == "__main__": raise SystemExit(main())