378 lines
14 KiB
Python
Executable File
378 lines
14 KiB
Python
Executable File
#!/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: <episode>/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())
|